what's a good way to save jax arrays?

hi all,

thanks for making such a cool library! I was wondering how I should go about saving jax arrays to file, since np.save isn’t yet implemented. Apologies if I missed something!

Rowan

[mattjj]

Thanks for the question!

JAX’s device-backed ndarray class, DeviceArray, is effectively a subclass of numpy.ndarray. You can turn it into a regular ndarray using numpy.asarray, like this:

import numpy as onp
numpy_array = onp.asarray(device_array)

DeviceArrays try to turn themselves into ndarrays automatically when appropriate, so you can also use Numpy’s save on them directly:

onp.save('array.npy', device_array)

We could also override __getstate__ so that they work with pickle. (We did that in an older version of JAX, but it looks like we lost that code at some point.)

I think we should probably bring numpy.save into jax.numpy, so that you don’t have to import raw NumPy (as onp here) to get the same saving behavior… Let’s leave this issue open until we either add jax.numpy.save or else figure out there’s some reason not to.

In the meantime, does import numpy as onp and onp.save work for you?

[rowanz]

Thanks for the response! Sorry, I’m still learning about the library and didn’t know that you could convert DeviceArrays into regular numpy arrays like that – that helps a lot and solves my problem (and will hopefully help others if they search for it) 😃