re: usage and typing of array I/O with NumPy exports

[nicholasjng]

Hey. I'm trying to skirt around the (for me unnecessary) dependency on tensorflow-datasets when loading MNIST.

For now, I came up with the following, which loads the train/test contents into uint8 arrays via frombuffer(): https://github.com/aai-institute/nnbench/blob/e291f71d1f7e69083c3f243a016359e9785523af/examples/mnist/mnist.py#L89-L122

I'd like to assert that what comes out here are actually JAX arrays, which I'm not sure is the case, since some of JAX's array I/O APIs are re-exports from NumPy (see e.g. jax.numpy.savez).

TL,DR: Is it correct to just go jnp.{load,frombuffer}() here, and type everything as a jax.Array, or do I need to coerce with asarray somewhere on the return values of jnp.load() or jnp.frombuffer()?

[jakevdp]

jax.numpy.frombuffer will always return JAX arrays. jax.numpy.load will return JAX arrays from .npy files, but not from .npz files. This is because the object returned by np.load when passed a .npz file is a non-trivial lazy view of the buffers on disk, and there was no easy way to hook into this in order to make it return JAX arrays. It's something that could probably be addressed with some effort if it were important, but nobody has told us before now that it's important.

Also, NumPy's array serialization formats (.npy, .npz) are not particularly well-suited for JAX, because they pre-suppose the particular set of dtypes that NumPy supports, and so there's no way to specify bfloat16 or other dtype extensions.

For your case, how hard is it to post-process everything with jnp.asarray?

[nicholasjng]

Thanks. It's not hard at all, no worries. I was more interested in the behavior of np.load() in particular.

I don't really want to give up the compression that I get with uint8 + .npz files, so I guess I can just tree_map the result before returning.