How to save an array in jitted function?

[drewjin]

I met with a problem dumping jax arrays in a complex jitted funtion.

I want to dump two jax arrays which is computed by the following function:

def fourier_embeddings(x: jnp.ndarray, dim: int) -> jnp.ndarray:
  w_key, b_key = jax.random.split(jax.random.PRNGKey(42))
  weight = jax.random.normal(w_key, shape=[dim])
  bias = jax.random.uniform(b_key, shape=[dim])
  return (
    jnp.cos(2 * jnp.pi * (x[..., None] * weight + bias)),
    weight, bias
  )

Following the community discussion:

#712
#3766

I wrote this function, hoping it will work:

def jax_save_with_jit(file: str, arr):
  # def save_to_file(a):
  #   jax.numpy.save(file, a)
  # hcb.id_tap(save_to_file, arr)
  if file == 'w':
    jax.debug.print('writing {} fourier_weight.npy', arr)
    host_callback.call(
      lambda arr: jnp.save('/home/user3/alphafold/model/fourier_weight.npy', arr), arr)
  elif file == 'b':
    jax.debug.print('writing {} fourier_bias.npy', arr)
    host_callback.call(
      lambda arr: jnp.save('/home/user3/alphafold/model/fourier_bias.npy', arr), arr)

The function call is as follows:

noise_embedding, weight, bias = fourier_embeddings(
  (1 / 4) * jnp.log(noise_level / SIGMA_DATA), dim=256
)
jax_save_with_jit('w', weight)
jax_save_with_jit('b', bias)

For the value type:

weight = Traced<ShapedArray(float32[256])>with<DynamicJaxprTrace(level=2/0)>
bias = Traced<ShapedArray(float32[256])>with<DynamicJaxprTrace(level=2/0)>

Unfortunately, it turns out that nothing was saved at all in the directed path...

[jakevdp]

Note that host_callback is deprecated; you probably want io_callback. For more information, see https://jax.readthedocs.io/en/latest/external-callbacks.html#exploring-io-callback

[drewjin]

Appreciate that!