Different behavior when loading .npy and .npz files using jnp.load

[Xin-yang-Liu]

I am loading data from a .npz file using jnp.load. However, it returns a numpy array instead of a jax device array. But when loading from .npy file, it works as I expected. I wonder if it is a BUG of jax or intended behavior.

To reproduce:

>>> import jax.numpy as jnp
>>> a=jnp.ones(5)
>>> jnp.savez('test',a=a)
>>> type(a)
<class 'jaxlib.xla_extension.ArrayImpl'>
>>> data=jnp.load('test.npz')
>>> type(data['a'])
<class 'numpy.ndarray'>
>>> jnp.save('test2',a)
>>> type(jnp.load('test2.npy'))
<class 'jaxlib.xla_extension.ArrayImpl'>

[jakevdp]

Thanks for the question. JAX's jnp.save and jnp.savez are just aliases of numpy's functions, and jnp.load is a light wrapper around numpy's load that only does special handling for the single-array case. Given that, the behavior you observe is expected.

It would be nice if we could make jnp.savez/jnp.load work seamlessly with JAX arrays, but I suspect that would be a bit of a project because it would involve re-implementing numpy's NpzFile and other structures from scratch. I think it's unlikely that that work will be done any time soon. In the meantime, I'd suggest that you anticipate that arrays loaded from npz files will be NumPy arrays rather than JAX arrays, and use jnp.asarray to convert them if necessary.

[Xin-yang-Liu]

Thanks! It would be great if you can mention this in the doc of jnp.load.