Why does io_callback in a custom JVP rule throw an error about unsupported JVP?

[mariogeiger]

I expected that by defining a custom jvp rule for my function, the io_callback used inside the custom rule wouldn’t need to support JVP itself. However, I’m getting the error: ValueError: IO callbacks do not support JVP. Could you explain why this happens and suggest how I might modify my code to call io_callback in a custom jvp?

import jax
from jax.experimental import io_callback


def _f(x):
    print(f"Running f with x={x}")
    return x


@jax.custom_jvp
def f(x):
    return x


@f.defjvp
def f_jvp(primals, tangents):
    (x,) = primals
    (dx,) = tangents
    io_callback(_f, x, jax.lax.stop_gradient(x))
    return x, dx


@jax.jit
def test(x):
    return jax.jvp(f, [x], [x])


test(1.0)

ValueError: IO callbacks do not support JVP.