{"record":{"id":"3cf54e43e359a594","repo":"jax-ml/jax","slug":"pure-callbacks-do-not-support-jvp-please-use-jax","errorCode":null,"errorMessage":"Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients.","messagePattern":"Pure callbacks do not support JVP\\. Please use `jax\\.custom_jvp` to use callbacks while taking gradients\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":117,"sourceCode":"pure_callback_p.def_impl(functools.partial(dispatch.apply_primitive,\n                                           pure_callback_p))\n\n\n@pure_callback_p.def_abstract_eval\ndef pure_callback_abstract_eval(\n    *avals,\n    callback: _FlatCallback,\n    result_avals,\n    sharding: Sharding | None,\n    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","sourceCodeStart":99,"sourceCodeEnd":135,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L99-L135","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Wrap the callback in @jax.custom_jvp defining a forward rule (often tangents of zeros or identity)","Use jax.lax.stop_gradient on the callback's inputs if gradients need not flow through it","Replace pure_callback with reimplemented JAX code when the derivative is known","For VJP-only needs, use @jax.custom_vjp instead"],"exampleFix":"# before\n@jax.jit\ndef f(x):\n    return pure_callback(scipy_fn, x, result_dtype=float)\njdx = jax.grad(f)(x)  # ValueError\n\n# after\n@jax.custom_jvp\ndef f(x):\n    return pure_callback(scipy_fn, x, result_dtype=float)\n@f.defjvp\ndef f_jvp(primals, tangents):\n    (x,), (xdot,) = primals, tangents\n    return f(x), jnp.zeros_like(x)  # or real derivative","handlingStrategy":"validation","validationCode":"import jax\n# Guard: differentiate only callback-free or custom_jvp-wrapped callables\ndef is_autodiff_safe(f, x):\n    try:\n        jax.jvp(f, (x,), (jax.tree.map(jnp.ones_like, x),))\n        return True\n    except ValueError as e:\n        return 'custom_jvp' not in str(e)","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept ValueError as e:\n    if 'do not support JVP' in str(e):\n        f = make_custom_jvp_version(f)  # user-defined wrapper\n    raise","preventionTips":["Define @jax.custom_jvp at function-definition time for anything wrapping pure_callback","Apply jax.lax.stop_gradient to non-differentiable host calls","Keep a registry mapping callback functions to their jvp/vjp rules"],"tags":["jax","autodiff","jvp","pure-callback","custom-jvp"],"backgroundTag":"jax-custom-jvp-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}