Add device argument for safetensors.flax.load_file

[mar-muel]

Feature request

Hey there - love this library! 👍

Any reason why the device argument is not valid (anymore?) for load_file for flax?

Also a bit confused, as it is listed as an argument in the docs 🤔 https://huggingface.co/docs/safetensors/main/en/api/flax#safetensors.flax.load_file

I'm using safetensors==0.4.1

Motivation

It's useful to have control over device placement during model load

Your contribution

Probably not...

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days.

[Narsil]

Thanks for the note, the docstring is outdated or bad copy pasted.

The reason for the argument not being here, is that Flax doesn't provide a way to create tensors directly on device (afaik),
meaning it's not going to yield any differences from loading on CPU then moving to whatever device.

Also I thought lazy tensors placements for flax was more idiomatic. How exactly do you move the tensors ?

[mar-muel]

@Narsil I later found out I can load my Flax msgpack models directly to CPU with this:

cpu_device = jax.devices('cpu')[0]
with jax.default_device(cpu_device):
    with open(msgpack_file, "rb") as state_f:
        state = from_bytes(cls, state_f.read())

In any case, I've now moved to saving my flax model in numpy format - which is what you get if you use jax.device_get():

>>> x = jnp.zeros((5,5))
>>> type(x)
<class 'jaxlib.xla_extension.ArrayImpl'>   # array is on device
>>> type(jax.device_get(x))
<class 'numpy.ndarray'>

Then to load the model

state = safetensors.numpy.load_file(st_file)   # np arrays on CPU
state = jax_utils.replicate(state)  # returns jax.numpy arrays replicated on default device

[Narsil]

Thanks for sharing your fix.