Saving large jax.arrays via jnp.savez or np.savez is slow - How to speed it up?

[isaac-tes]

As far as I understand, both np.savez and jnp.savez both convert jax.array() into np.array() via np.asarray() under the hood.
My computation greatly improves upon using jax, but the bottleneck is really the simple saving routine.
So the entire speed advantage I get from jax.numpy.linalg routines is gone when I need to save the larger subsets of the data, e.g., the squared eigenvectors of some eigendecomposition.

Is there anything I am missing how to save larger jax.arrays() more efficiently?

I am using jax entirely on CPUs if that is of help.
If you need additional information or some simple code snippets showing the slow saving routines, please let me know.
Any help is greatly appreciated, thanks a lot.

[jakevdp]

jnp.savez is just a thin wrapper around np.savez, and both work by converting a JAX array to NumPy via an implicit device_get. and then calling np.savez. If you're using JAX entirely on CPU, then the array conversion should be virtually free (the memory buffer will be shared between JAX and NumPy) so any bottleneck you're seeing is purely on the NumPy side of things.

Is there anything I am missing how to save larger jax.arrays() more efficiently?

Can you say more about what your ultimate goal is? If your hope is to generate .npz files containing your array content, I don't think there's any improvement to be had. But if, say, you're doing more general checkpointing within JAX workflows, something like orbax-checkpoint might give you better performance.