jax.debug.impure_callback

[pchanial]

In order to polish some informative messages and because jax.debug.callback is traced, I resorted to jax.experimental.io_callback to print a conditional message:

def callback(solution):
    niter = solution['niter']
    if niter < nitermax:
        print(f'Converged in {niter} iterations')
    else:
        print(f'Did not converge in {niter} iteration')
    return jax.array(0, dtype=np.int8)

which is then called in a staged out function with:

_ = jax.experimental.io_callback(callback, jax.ShapeDtypeStruct((), np.int8), solution)

It would be cleaner to just have to write it without the trick of returning a useless scalar with a function jax.debug.impure_callback(callback, solution)

[jakevdp]

It sounds like what you're asking for already exists in jax.debug.callback. Can you say more about why this isn't sufficient for your use-case?

[pchanial]

My apologies, I cannot reproduce the problem I encountered. I will reopen the discussion if I see again the issue.