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

[usernamesjh]

Dear developers,

Hello! I encountered a difficult problem while using Netket, 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?

[PhilipVinc]

This is not a netket-specific question, but rather a jax specific one. You should ask it other at https://github.com/jax-ml/jax .
This is a common question in jax.

Regardless, the issue here is that you have a network and it is Jitted, so you cannot insert arbitrary functions.
The solution would be to use a jax.pure_callback or more appropriately a jax.experimental.io_Callback if you want to save to file.

If you want to print the content of an array inside a jitted function, instead, you should not do print(x) but you should use jax.debug.print and do something like

jax.debug.print("my array {}", x)

[usernamesjh]

Thank you very much for your reply! Your advice has been very inspiring to me. I will try it as soon as possible according to what you said, and I hope to solve the problem smoothly.