Quick question on implementation.

[miguelamendez]

Hi guys,
Just a quick note/question for the way the library is implemented.
I realised that jax natively has safe and loading trhoguh out jax.numpy.sava/savez: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.save.html.
Also, safetensor library has a module safetensor.numpy wich lets you save and load natively numpy and also jax.numpy arrays:
https://huggingface.co/docs/safetensors/api/numpy
from safetensors.numpy import save_file
import jax.numpy as jnp
tensors = {"embedding": jnp.zeros((4, 8)), "attention": jnp.zeros((4,4))}
save_file(tensors, "model.safetensors")
print(tensors)
from safetensors.numpy import load
file_path = "./model.safetensors"
with open(file_path, "rb") as f:
data = f.read()
loaded = load(data)
print(loaded)

Does this library adds any increase of performance over using the safetensors native numpy save and load functions?, If so shouldn't be merged to safe-tensors repo directly?

[alvarobartt]

Hi @miguelamendez so the idea behind safetensors is that it's another format for storing tensors in opposition to pickle which is not safe. This doesn't mean that jax.numpy.save, jax.numpy.savez, or any other existing serializing format is not recommended.

The idea behind safejax is to use safetensors to serialize the whole tree, which decomposes the trees in a Python dictionary by joining the keys, so that the original tree is flattened. This is done because safetensors doesn't support the storing of complex structures/trees e.g. FrozenDict in flax. More information on the latter is at huggingface/safetensors#138

So on, the following doesn't work with plain safetensors:

from safetensors.jax import save_file

tensors = {"dense0": {"weight": jnp.ones((1, 10)), "bias": jnp.ones((10))}}
save_file(tensors, "tensors.safetensors")

While using safejax, the structure is nested and then serialized with safetensors:

from safejax import serialize

tensors = {"dense0": {"weight": jnp.ones((1, 10)), "bias": jnp.ones((10))}}
serialize(tensors, filename="tensors.safetensors")

More information on why safejax was created at https://github.com/alvarobartt/safejax/tree/main#-why-safejax

For extra context, safejax works in the case above while safetensors doesn't because safetensors just supports level-1 Python dictionaries, so safejax solves that by flattening the dictionaries to be of level-1 too. In the example above, {"dense0": {"weight": jnp.ones((1, 10)), "bias": jnp.ones((10))}} ends up being {"dense0.weight": jnp.ones((1, 10)), "dense0.bias": jnp.ones((10))} so that can indeed be serialized with safetensors.

[miguelamendez]

Thanks for clarification!