🔐 Serialize JAX, Flax, Haiku, or Objax model params with safetensors

safejax is a Python package to serialize JAX, Flax, Haiku, or Objax model params using safetensors as the tensor storage format, instead of relying on pickle. For more details on why safetensors is safer than pickle please check huggingface/safetensors.

Note that safejax supports the serialization of jax, flax, dm-haiku, and objax model parameters and has been tested with all those frameworks, but there may be some cases where it does not work as expected, as this is still in an early development phase, so please if you have any feedback or bug reports, open an issue at safejax/issues.

---

Last update: 2023-01-19
Created: 2023-01-19