From traced array to numpy array

[mattiadutto]

I'm using Jax for a Federated Learning project, I want to check a few values inside different traced arrays and compere with values inside numpy arrays. So I have the following questions:

• There is any way to move from traced array to a numpy array?
• Or a way to get access to the real value inside a cell of the traced array?

The best that I found is array.at[index].get() but unfortunately it return a traced array.

Thanks.
Mattia

[jakevdp]

Edit: there is now a FAQ entry on this topic: https://jax.readthedocs.io/en/latest/faq.html#how-can-i-convert-a-jax-tracer-to-a-numpy-array

---

Thanks for the question! You can print traced values at runtime using jax.debug.print. Note however that the syntax of debug.print is different than that of normal print statements: the first argument must be a format string. It might look something like this:

import jax
import jax.numpy as jnp

@jax.jit
def f(x):
  print("traced:", x[0])
  jax.debug.print("debug: {}", x[0])

f(jnp.arange(10))

traced: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
debug: 0

[mattiadutto]

Thanks for the quick answer!
So i can use jax.debug.print() only inside a compiled jit function.

For the part how to move to a Numpy array do you have any suggestions?

[jakevdp]

Sorry - I think I missed that part of the question.

You cannot convert a tracer into a numpy array. But if you want to execute numpy code in the midst of a JAX execution, you can do so via a callback to the host using jax.pure_callback. It's fairly efficient on CPU, where the host-to-device memory transfers can usually be done in a zero-copy manner, but be aware that on GPU the device transfer overhead and computational lock can cause severe performance hits.

There's not great documentation on this yet, but there's a short demo in this answer: #12180 (comment)

[mattiadutto]

Don't worry, thanks!