{"record":{"id":"0e72693bb8a776fc","repo":"jax-ml/jax","slug":"pure-callbacks-do-not-support-transpose-please-us","errorCode":null,"errorMessage":"Pure callbacks do not support transpose. Please use `jax.custom_vjp` to use callbacks while taking gradients.","messagePattern":"Pure 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/callback.py","lineNumber":127,"sourceCode":"    vmap_method: str | None,\n):\n  del avals, callback, sharding, vmap_method\n  return result_avals\n\n\ndef pure_callback_jvp_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Pure callbacks do not support JVP. \"\n      \"Please use `jax.custom_jvp` to use callbacks while taking gradients.\")\n\n\nad.primitive_jvps[pure_callback_p] = pure_callback_jvp_rule\n\n\ndef pure_callback_transpose_rule(*args, **kwargs):\n  del args, kwargs\n  raise ValueError(\n      \"Pure callbacks do not support transpose. \"\n      \"Please use `jax.custom_vjp` to use callbacks while taking gradients.\")\n\nad.primitive_transposes[pure_callback_p] = pure_callback_transpose_rule\n\n\nbatching.primitive_batchers[pure_callback_p] = functools.partial(\n    ffi.ffi_batching_rule, pure_callback_p\n)\n\ndef _get_sdy_array_list_for_callbacks(avals: Sequence[core.ShapedArray]) -> SdyArrayList:\n  \"\"\"Returns an SdyArrayList with `max(1, len(avals))` replicated shardings.\"\"\"\n  ndims = [0]\n  if avals:\n    ndims = [x.ndim for x in avals if isinstance(x, core.ShapedArray)]\n  return SdyArrayList(tuple(\n      SdyArray(\n          mesh_shape=(),","sourceCodeStart":109,"sourceCodeEnd":145,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L109-L145","documentation":"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.","triggerScenarios":"jax.grad / jax.vjp / jax.linear_transpose over a function whose trace includes jax.pure_callback or jax.io_callback.","commonSituations":"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.","solutions":["Apply jax.lax.stop_gradient to the pure_callback inputs/outputs so transpose never reaches it","Define @jax.custom_vjp with an explicit bwd rule (e.g. pass-through or zeros)","Reimplement the callback'd math in native JAX ops so autodiff works automatically"],"exampleFix":"# before\ndef loss(x):\n    return jnp.sum(pure_callback(fn, x, result_dtype=float))\njax.grad(loss)(x)  # ValueError\n\n# after\nfrom jax import lax\nndef loss(x):\n    return jnp.sum(pure_callback(fn, lax.stop_gradient(x), result_dtype=float))\njax.grad(loss)(x)  # ok, grad is zero through callback","handlingStrategy":"validation","validationCode":"import jax\ntry:\n    jax.vjp(f, x)\n    safe = True\nexcept ValueError:\n    safe = False  # wrap f with custom_vjp before calling grad","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        f = with_custom_vjp(f)\n    raise","preventionTips":["Use jax.lax.stop_gradient on all pure_callback inputs by default","Centralize host-interop functions behind custom_vjp wrappers","Unit-test grad on every loss containing callbacks"],"tags":["jax","autodiff","vjp","transpose","pure-callback","custom-vjp"],"backgroundTag":"jax-custom-vjp-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}