jax-ml/jax · error · ValueError

Buffer callbacks do not support transpose. Please use `jax.c

Error message

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

What it means

buffer_callback is a JAX primitive that ships raw buffers to a Python callback; JAX registers a transpose rule for it that unconditionally raises, because the JAX transform system needs a transpose rule for every primitive when computing VJPs. Taking a gradient (grad, vjp, or any transformation that transposes) through a buffer_callback has no defined mathematical meaning, so JAX refuses and points you at jax.custom_vjp to define the backward pass yourself.

Source

Thrown at jax/_src/buffer_callback.py:214

    has_side_effect: bool,
    **_,
):
  del args
  effects = {_BufferCallbackEffect} if has_side_effect else core.no_effects
  return result_avals, effects


def _buffer_callback_jvp_rule(*args, **kwargs):
  del args, kwargs
  raise ValueError(
      "Buffer callbacks do not support JVP. "
      "Please use `jax.custom_jvp` to use callbacks while taking gradients.")
ad.primitive_jvps[buffer_callback_p] = _buffer_callback_jvp_rule


def _buffer_callback_transpose_rule(*args, **kwargs):
  del args, kwargs
  raise ValueError(
      "Buffer callbacks do not support transpose. "
      "Please use `jax.custom_vjp` to use callbacks while taking gradients.")
ad.primitive_transposes[buffer_callback_p] = _buffer_callback_transpose_rule

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


def _buffer_callback_lowering(
    ctx: mlir.LoweringRuleContext,
    *args: Any,
    callback,
    in_tree: Any,
    out_tree: Any,
    has_side_effect: bool,
    input_output_aliases: Sequence[tuple[int, int]],
    command_buffer_compatible: bool,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the callback-containing function with @jax.custom_vjp and define an explicit bwd rule (usually passing zeros/matching cotangents through)
2. Alternatively use jax.custom_jvp to define the forward-mode rule if only JVPs are needed
3. Move the callback outside the differentiated region (detach/stop_gradient the inputs, e.g. jax.lax.stop_gradient, before the callback)
4. Use pure_callback/io_callback which raise the analogous error earlier with clearer guidance

Example fix

// before
def loss(x):
  buf = buffer_callback(x)  # inside grad
  return jnp.sum(x)
jdx = jax.grad(loss)(x)  # ValueError

// after
@jax.custom_vjp
def loss(x):
  return _loss_fwd_impl(x)
def loss_fwd(x):
  return _loss_fwd_impl(x), None
def loss_bwd(_, g):
  return g  # callback is observation-only
loss.defvjp(loss_fwd, loss_bwd)
jdx = jax.grad(loss)(x)

import jax
# Before differentiating, ensure no buffer/pure callbacks sit in the traced function
# simplest guard: stop_gradient inputs destined for the callback
x_sg = jax.lax.stop_gradient(x)  # then feed x_sg to buffer_callback

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

When it happens

Trigger: Calling jax.grad, jax.vjp, jax.linear_transpose, or any differential operator on a function whose trace reaches buffer_callback (e.g. emitted by experimental buffer_callback APIs or debugger/DAP buffer inspection inside a differentiable computation).

Common situations: Using experimental debugging/buffer-inspection callbacks inside a training loss; wrapping visualization or snapshotting code that captures activations inside grad-traced functions; assuming callbacks are transparent to autodiff.

Related errors

• buffer_callback callback must not return any values.
• Called {name} with a float0 array. float0s do not support an
• Cannot transpose with respect to sparse indices
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t

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