jax-ml/jax · error · ValueError

Pure callbacks do not support transpose. Please use `jax.cus

Error message

Pure callbacks do not support transpose. Please use `jax.custom_vjp` to use callbacks while taking gradients.

What it means

Transposition (the backward half of VJP) requires a transpose rule per primitive; pure_callback is opaque so its transpose rule always raises, telling you to define the backward pass with jax.custom_vjp. Unlike JVP, there is no safe default because the callback's linear structure is unknown to JAX.

Source

Thrown at jax/_src/callback.py:127

    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
)

def _get_sdy_array_list_for_callbacks(avals: Sequence[core.ShapedArray]) -> SdyArrayList:
  """Returns an SdyArrayList with `max(1, len(avals))` replicated shardings."""
  ndims = [0]
  if avals:
    ndims = [x.ndim for x in avals if isinstance(x, core.ShapedArray)]
  return SdyArrayList(tuple(
      SdyArray(
          mesh_shape=(),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply jax.lax.stop_gradient to the pure_callback inputs/outputs so transpose never reaches it
2. Define @jax.custom_vjp with an explicit bwd rule (e.g. pass-through or zeros)
3. Reimplement the callback'd math in native JAX ops so autodiff works automatically

Example fix

# before
def loss(x):
    return jnp.sum(pure_callback(fn, x, result_dtype=float))
jax.grad(loss)(x)  # ValueError

# after
from jax import lax
ndef loss(x):
    return jnp.sum(pure_callback(fn, lax.stop_gradient(x), result_dtype=float))
jax.grad(loss)(x)  # ok, grad is zero through callback

import jax
try:
    jax.vjp(f, x)
    safe = True
except ValueError:
    safe = False  # wrap f with custom_vjp before calling grad

try:
    jax.grad(f)(x)
except ValueError as e:
    if 'do not support transpose' in str(e) and 'custom_vjp' in str(e):
        f = with_custom_vjp(f)
    raise

When it happens

Trigger: jax.grad / jax.vjp / jax.linear_transpose over a function whose trace includes jax.pure_callback or jax.io_callback.

Common situations: Calling grad on losses that embed SciPy/BLAS/Numerical-Routines via pure_callback; using pure_callback for I/O or logging inside training graphs; assuming lax.stop_gradient was applied but forgetting it on one path.

Related errors

• transpose with implicit broadcasting of unshaped values. Got
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• structure of the differentiated function {jaxpr.debug_info.f
• Pure callbacks do not support JVP. Please use `jax.custom_jv

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