Conversion to NumPy array is slow

[allentran]

I have a really odd bug that I can't sort out...

I have a JAX float32-array of shape (2, 7, 37, 4, 3, 2) that takes 4 hours to convert to a NumPy array (i.e np.array(jax_array)). If I do not call np.array and just return the JAX array the function call takes seconds to return. Initially I thought I'd see what the max and minimum are but jnp.max() and jnp.min() also take ~4 hours.

I cannot reproduce it outside of the codebase with a similar shaped randomized matrix (conversions take < 1 second). The array is formed within an experimental loop replicating a markov chain but otherwise unexceptional.

Does anyone have any ideas about how to debug this?

[hawkinsp]

Is it that the array conversion is slow, or that the Jax array is actually a future that isn't ready yet?

Does .block_until_ready() also block for a long time on the JAX array? If so, it's most likely that JAX is computing the value in the background and it simply isn't ready for a long time.

[allentran]

Aha! That was it, thank you for this, it was driving me insane. Although now I have the bigger issue of shaving 4 hours into something manageable.