jax-ml/jax · error · ValueError

Pure callbacks do not support JVP. Please use `jax.custom_jv

Error message

Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients.

What it means

JVP differentiation requires a JVP rule for every primitive in the traced function. pure_callback treats the callback as an opaque black box with no known derivative, so its registered JVP rule unconditionally raises ValueError and directs you to jax.custom_jvp, where you define the tangent rule yourself (typically zero or a pass-through).

Source

Thrown at jax/_src/callback.py:117

pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,
                                           pure_callback_p))


@pure_callback_p.def_abstract_eval
def pure_callback_abstract_eval(
    *avals,
    callback: _FlatCallback,
    result_avals,
    sharding: Sharding | None,
    vmap_method: str | None,
):
  del avals, callback, sharding, vmap_method
  return result_avals


def pure_callback_jvp_rule(*args, **kwargs):
  del args, kwargs
  raise ValueError(
      "Pure callbacks do not support JVP. "
      "Please use `jax.custom_jvp` to use callbacks while taking gradients.")


ad.primitive_jvps[pure_callback_p] = pure_callback_jvp_rule


def pure_callback_transpose_rule(*args, **kwargs):
  del args, kwargs
  raise ValueError(
      "Pure callbacks do not support transpose. "
      "Please use `jax.custom_vjp` to use callbacks while taking gradients.")

ad.primitive_transposes[pure_callback_p] = pure_callback_transpose_rule


batching.primitive_batchers[pure_callback_p] = functools.partial(
    ffi.ffi_batching_rule, pure_callback_p

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the callback in @jax.custom_jvp defining a forward rule (often tangents of zeros or identity)
2. Use jax.lax.stop_gradient on the callback's inputs if gradients need not flow through it
3. Replace pure_callback with reimplemented JAX code when the derivative is known
4. For VJP-only needs, use @jax.custom_vjp instead

Example fix

# before
@jax.jit
def f(x):
    return pure_callback(scipy_fn, x, result_dtype=float)
jdx = jax.grad(f)(x)  # ValueError

# after
@jax.custom_jvp
def f(x):
    return pure_callback(scipy_fn, x, result_dtype=float)
@f.defjvp
def f_jvp(primals, tangents):
    (x,), (xdot,) = primals, tangents
    return f(x), jnp.zeros_like(x)  # or real derivative

import jax
# Guard: differentiate only callback-free or custom_jvp-wrapped callables
def is_autodiff_safe(f, x):
    try:
        jax.jvp(f, (x,), (jax.tree.map(jnp.ones_like, x),))
        return True
    except ValueError as e:
        return 'custom_jvp' not in str(e)

try:
    jax.grad(f)(x)
except ValueError as e:
    if 'do not support JVP' in str(e):
        f = make_custom_jvp_version(f)  # user-defined wrapper
    raise

When it happens

Trigger: Calling jax.jvp, or jax.grad/jax.value_and_grad (which use JVP internally via VJP construction) on a function containing jax.pure_callback; also jax.checkpoint/jacfwd traces.

Common situations: Mixing numerical libraries (SciPy, NumPy RNG, custom kernels) into differentiable JAX code via pure_callback and then calling grad; porting PyTorch code that used custom autograd functions; using pure_callback for data-dependent logic inside a loss.

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim
• Pure callbacks do not support transpose. Please use `jax.cus

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3cf54e43e359a594. Report an issue: GitHub.