Conditional saving/storing in a jitted function

[andremfreitas]

Hi,

What is the optimal approach when we want to save some result every, for example, 10 steps, of a numerical solver ? jax.lax.cond seems not to be ideal because the true and false branches need to return the same type.

PS: This numerical solver is called iteratively via a jax.lax.fori_loop or a jax.lax.scan

Thanks,
Andre

[jakevdp]

Saving a value is a side-effect, so you cannot do it directly in a jit-compiled function. But you should be able to use jax.experimental.io_callback for this; see a discussion at External Callbacks in JAX.

[andremfreitas]

Thanks! This is exactly what I was looking for :)