Saving intermediate values in a file

[simlaharma]

Hello,

I have a JAX function that I use during training a DNN model. I want to debug it by checking the values before and after the function.
So I have been looking for a solution to save the intermediate values during training. The only solution I came across was to use jax.debug.print. However, this only prints and I cannot print huge arrays. Is there a way similar to jax.numpy.save function but for traced arrays?

[JiaYaobo]

jax.experimental.io_callback?

def save_array(array):
    np.save('array.npy', array)
    print('saved array')

@jax.jit
def f(x):
    io_callback(save_array, None, x)
    return jax.numpy.sin(x)

[jakevdp]

This is the answer. For more information on callbacks and why they're necessary, see https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html.