saving jax tensor in jitted function

[jwnys]

Opening a new issue for my comment in #712

I need to save my tensor inside a jitted function. jax.numpy.save behaves the same as numpy.save and throws a conversion error (it's trying to convert to onp.ndarray) "The numpy.ndarray conversion method array() was called on the JAX Tracer object". Is there an alternative or solution to save tensors?

[hawkinsp]

(I converted your issue to a discussion, because a discussion seemed more appropriate.)

Yes, this is working as intended. Inside a jit block, you have Tracer objects, which are symbolic objects that stand in for arrays. You cannot save the value of one because it is a symbolic value that stands in for a set of concrete values, not a particular concrete value itself.

There are a number of solutions to this problem:
a) don't use jit while debugging. This means you have concrete values once again, as you would expect. This is certainly the first thing I usually try because it means all the usual debugging tools work: saving, printing, etc.

You mentioned that you had tried this and it was too slow. It may be worth looking into this. Can you say more about the performance differences and possibly share a reproduction? Perhaps there's something we should look into here.

b) you could plumb the values you want to save as extra outputs from the jit-compiled function. This will work, but may be awkward to do.

c) there's new, experimental host_callback support that does something fairly similar to what you want
https://github.com/google/jax/blob/master/jax/experimental/host_callback.py

It's specifically intended for making logging and debugging easier.

Does this answer your question?

[jwnys]

Thanks for the pointers!

a) to give some context: I'm using stax in combination with sampling algos. Every n'th step, I'd like to output tensors to see its evolution (this is not for debugging, but to monitor the evolution). Some code reproduction is non-trivial since it uses external libraries (https://github.com/netket/netket). Essentially, even if I can output the tensors for an iteration step, I need these during the whole optimization process, which takes forever when not jitted.

b) this indeed seems like an option, but it's quite a big change that requires adjusting other packages (to be honest, this seems like something that might lead me to avoid jax, but up to this issue, it can do everything that I need and have been looking for). Context: the external package I'm using just takes a neural network from stax and jits it. It's a nice flow that I hope would work in general.

c) the host_callback seems like something I need indeed, thanks for pointing out this reference!!! I didn't get it to work though, I was expecting something like this to be sufficient (based on the logging examples in host_callback_tests.py):

from jax.experimental import host_callback as hcb

def jax_save(file, arr):
    def save_to_file(a):
        jax.numpy.save(file, a)
    hcb.id_tap(save_to_file, arr)

Yet, this doesn't do anything (no crashing, no warnings,.., just no output). Maybe this is too naive?

Thanks!

[jekbradbury]

I think something like that should work; can you try debugging in steps from id_print to id_tap with printing to id_tap with saving?

[jayati-naik]

Hi, I'm also stuck at same issue. Did jax_save function worked ? Please confirm.

[sshleifer]

any update?

[yashk2810]

I would recommend using io_callback here: https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html#exploring-jax-experimental-io-callback