serialization.from_state_dict does not restore to jax.Arrays

[PhilipVinc]

Deserialising a single jax.Array with flax.serialization does not reconstruct jax.Arraycorrectly even when they are present in the target pytree.

See

import jax
import jax.numpy as jnp
import flax
from flax import serialization

pt = {'a':jnp.ones((3,4))}

bdata = serialization.to_bytes(pt)
pt_loaded = serialization.from_bytes(pt, bdata)

print("Standard Deserialization  output type:", type(pt_loaded['a']))
> Standard Deserialization output type: <class 'numpy.ndarray'>

This is because there is no rule for how to deal with jax.Arrays.

To fix it, one needs to register this rule

serialization.register_serialization_state(
  type(pt['a']),
  lambda x: x,
  lambda x, sd: jax.numpy.asarray(sd, dtype=x.dtype),
  override=True
)
bdata = serialization.to_bytes(pt)
pt_loaded = serialization.from_bytes(pt, bdata)

print("type Deserialization      output type:", type(pt_loaded['a']))
> type Deserialization      output type: <class 'jaxlib.xla_extension.ArrayImpl'>

Would it be possible to get this inside of flax itself?

[IvyZX]

Thanks for the suggestion.
One concern we may have is that JAX sets its array type as private (jax._src.array.ArrayImpl) to keep the flexibility to modify its APIs, especially because JAX arrays now cover a bunch of device-specific metadata (like shardings, memory placement, etc).
On the other hand, flax.serialization treats an array only by its values, and no metadata like device shardings etc will be preserved when the bytes get deserialized later. Thus it fits better with being a NumPy array simply sitting in the host CPU memory.
It isn't hard to convert it into a naive unsharded JAX array - just jax.tree.map(jnp.array, pt_loaded) is enough. But it might be better to leave to the user to fill out its jax array metadata.