JAX numpy to Numpy is CRAZILY SLOW

[LIANHAN3920]

I wrote an optimization function using jax, and used jax.vmap to run several optimization. The result is simply a (N, ) jax.numpy array. But when I try to convert it into a numpy array as python float, it takes fooooooooorever.

The time used to complete over 100 optimization is just several seconds, but to convert it into python float, it takes over 1 minuets!

I tried several methods, they are all slow.

np_as_array = np.asarray(test_1, dtype='float32') Time for using numpy.asarray: 58.892239799999516

np_array = np.array(test_2, dtype=float) #Time for using numpy.array: 69.0666779000021

jnp_array = test_3.__array__() #Time for using jnp.array: 61.6944502999977

jnp_item = [i.item() for i in test_4] #Time for using jnp.item(): 63.73181040000054

Did I do anything wrong or is there a faster way. Can not believe it :(
And they are all on CPU

[pearu]

In general, the np.asarray approach is preferred as it returns a view of the jax array (provided that dtypes match and jax array is on CPU) which is expected to be almost instant operation.

What is the size of these arrays and how much RAM is available in your system? Also, what is the dtype of the JAX arrays?

[hawkinsp]

Please read https://docs.jax.dev/en/latest/async_dispatch.html

Are you sure you are timing the conversion itself, as opposed to the computation? Try adding a jax.block_until_ready(test_1) before calling np.asarray.

[LIANHAN3920]

Thanks! I think that is the reason!