{"record":{"id":"ea74b35b08407e5b","repo":"jax-ml/jax","slug":"buffer-callbacks-do-not-support-transpose-please","errorCode":null,"errorMessage":"Buffer callbacks do not support transpose. Please use `jax.custom_vjp` to use callbacks while taking gradients.","messagePattern":"Buffer callbacks do not support transpose\\. Please use `jax\\.custom_vjp` to use callbacks while taking gradients\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/buffer_callback.py","lineNumber":214,"sourceCode":"    has_side_effect: bool,\n    **_,\n):\n  del args\n  effects = {_BufferCallbackEffect} if has_side_effect else core.no_effects\n  return result_avals, effects\n\n\ndef _buffer_callback_jvp_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Buffer callbacks do not support JVP. \"\n      \"Please use `jax.custom_jvp` to use callbacks while taking gradients.\")\nad.primitive_jvps[buffer_callback_p] = _buffer_callback_jvp_rule\n\n\ndef _buffer_callback_transpose_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Buffer callbacks do not support transpose. \"\n      \"Please use `jax.custom_vjp` to use callbacks while taking gradients.\")\nad.primitive_transposes[buffer_callback_p] = _buffer_callback_transpose_rule\n\nbatching.primitive_batchers[buffer_callback_p] = functools.partial(\n    ffi.ffi_batching_rule, buffer_callback_p\n)\n\n\ndef _buffer_callback_lowering(\n    ctx: mlir.LoweringRuleContext,\n    *args: Any,\n    callback,\n    in_tree: Any,\n    out_tree: Any,\n    has_side_effect: bool,\n    input_output_aliases: Sequence[tuple[int, int]],\n    command_buffer_compatible: bool,","sourceCodeStart":196,"sourceCodeEnd":232,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/buffer_callback.py#L196-L232","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Wrap the callback-containing function with @jax.custom_vjp and define an explicit bwd rule (usually passing zeros/matching cotangents through)","Alternatively use jax.custom_jvp to define the forward-mode rule if only JVPs are needed","Move the callback outside the differentiated region (detach/stop_gradient the inputs, e.g. jax.lax.stop_gradient, before the callback)","Use pure_callback/io_callback which raise the analogous error earlier with clearer guidance"],"exampleFix":"// before\ndef loss(x):\n  buf = buffer_callback(x)  # inside grad\n  return jnp.sum(x)\njdx = jax.grad(loss)(x)  # ValueError\n\n// after\n@jax.custom_vjp\ndef loss(x):\n  return _loss_fwd_impl(x)\ndef loss_fwd(x):\n  return _loss_fwd_impl(x), None\ndef loss_bwd(_, g):\n  return g  # callback is observation-only\nloss.defvjp(loss_fwd, loss_bwd)\njdx = jax.grad(loss)(x)","handlingStrategy":"validation","validationCode":"import jax\n# Before differentiating, ensure no buffer/pure callbacks sit in the traced function\n# simplest guard: stop_gradient inputs destined for the callback\nx_sg = jax.lax.stop_gradient(x)  # then feed x_sg to buffer_callback","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept ValueError as e:\n    if 'do not support transpose' in str(e) and 'custom_vjp' in str(e):\n        # wrap f with jax.custom_vjp and retry\n        ...\n    raise","preventionTips":["Never place observation callbacks inside grad-traced regions","Wrap callbacks with jax.custom_vvp/custom_vjp at module definition time, not after hitting the error","Use jax.lax.stop_gradient on callback inputs by convention when gradients are meaningless"],"tags":["jax","autodiff","gradient","callback","buffer-callback"],"backgroundTag":"jax-custom-vjp-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}