How to save Jax.interpreters.ad.JVPTracer as numpy array?

[ghost]

In the latest JAX, we can print values using jax.debug.print function (https://jax.readthedocs.io/en/latest/debugging/print_breakpoint.html). However, how to save those values as a numpy array? Thank you for your time.

[sharadmv]

You can use jax.debug.callback which will provide the NumPy arrays to a Python callback. You can then save them however you want.