How can we export the data generated during the calculation process?

[kaho-sue]

Dear developers,

Hello! I encountered a difficult problem while using jax, and I hope to get your help here. I have tried various methods, but the problem still persists, so I decided to seek advice from you.
We aim to export the data generated throughout the calculation process. However, I encounter certain errors when attempting to do so.

In this photo, my intention was to print(x), yet we are only able to print the data type instead.

I searched for a solution on Google. It suggests using jax.config.update('jax_disable_jit', True). After that, I can print the specific data.

However, we are unable to save this data.

So, what's the proper way to save this data?
Would the solution be to use a jax.pure_callback or more appropriately a jax.experimental.io_Callback if you want to save to file?

[jakevdp]

Hi - thanks for the question! The issue here is that print occurs during tracing (see JAX key concepts: Tracing for some background), so when the print statement executes your data has not yet been passed to the program.

The solution would be to use some sort of callback (See External Callbacks for an intro to these); likely either jax.debug.callback or jax.experimental.io_callback depending on your needs.

The simplest approach would be to use jax.debug.print directly; for example replace print(x) with jax.debug.print("{}", x).