{"record":{"id":"e1c3e9c8244fc1f5","repo":"jax-ml/jax","slug":"effects-not-supported-in-custom-jvp-disallowed","errorCode":null,"errorMessage":"Effects not supported in `custom_jvp`: {disallowed_effects}","messagePattern":"Effects not supported in `custom_jvp`: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":437,"sourceCode":"    out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents)\n    out_primals, nz_out_tangents = split_list(out, [len(out_zeros)])\n    nz_out_tangents_ = iter(nz_out_tangents)\n    out_tangents = [SymbolicZero(core.typeof(p).to_tangent_aval())\n                    if z else next(nz_out_tangents_)\n                    for p, z in zip(out_primals, out_zeros)]\n    assert next(nz_out_tangents_, None) is None\n    return [*out_primals, *out_tangents]\n  return lu.wrap_init(jvp, debug_info=jvp_jaxpr_fun.debug_info)\n\ncustom_jvp_call_p = CustomJVPCallPrimitive('custom_jvp_call')\n\ndef _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_fun,\n                               num_consts, symbolic_zeros):\n  # TODO(mattjj): could do more checking here...\n  del in_avals, jvp_jaxpr_fun, num_consts\n  disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(call_jaxpr.effects)\n  if disallowed_effects:\n    raise NotImplementedError(\n        f'Effects not supported in `custom_jvp`: {disallowed_effects}')\n  return call_jaxpr.out_avals, core.positional_effects(call_jaxpr)\ncore.custom_typechecks[custom_jvp_call_p] = _custom_jvp_call_typecheck\n\ndef _custom_jvp_vjp_call_lowering(ctx: mlir.LoweringRuleContext, *args,\n                                  call_jaxpr: core.Jaxpr, **_):\n  consts = mlir.ir_consts(\n      call_jaxpr.consts, [v.aval for v in call_jaxpr.constvars])\n  out, tokens = mlir.jaxpr_subcomp(ctx.module_context, call_jaxpr,\n                                   ctx.name_stack, ctx.tokens_in, consts,\n                                   *args, dim_var_values=ctx.dim_var_values,\n                                   const_lowering=ctx.const_lowering,\n                                   outer_traceback=ctx.traceback)\n  ctx.set_tokens_out(tokens)\n  return out\nmlir.register_lowering(custom_jvp_call_p, _custom_jvp_vjp_call_lowering)\n\ndef _custom_jvp_call_transpose_fancy(params, jaxpr, args, ct, _):","sourceCodeStart":419,"sourceCodeEnd":455,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L419-L455","documentation":"Effects (e.g. ordered effects like state or RNG primitives) inside the body of a @jax.custom_jvp function are not supported — the typecheck for the custom_jvp_call primitive filters the jaxpr's effects against the allowlist for custom derivatives and raises NotImplementedError otherwise.","triggerScenarios":"Placing effectful operations (random draws via stateful RNG APIs, mutable state updates, host callbacks) inside a function decorated with @jax.custom_jvp and tracing it (e.g. under jit/grad).","commonSituations":"Using new-style random or haiku/flax state inside custom derivative kernels; adding a print/halt/debug callback for debugging inside a custom_jvp function.","solutions":["Move effectful operations (randomness, state reads/writes) outside the custom_jvp function and pass their results in as arguments","Use pure jax.random PRNG keys passed as arguments instead of stateful RNG","Remove debug callbacks/halts from the decorated body"],"exampleFix":"# before\n@jax.custom_jvp\ndef f(x):\n    key = next_key()  # effectful\n    return x + jax.random.normal(key, x.shape)\n# after\ndef make(x, key):\n    return x + jax.random.normal(key, x.shape)\n@jax.custom_jvp\ndef f(x, noise):\n    return x + noise\n# caller draws noise with an explicit key beforehand","handlingStrategy":"validation","validationCode":"jaxpr = jax.make_jaxpr(f)(*sample_args)\nassert not jaxpr.effects, f'effects present: {jaxpr.effects}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep custom_jvp bodies pure; hoist RNG/state outside","Pass PRNG keys as traced arguments","Use jax.make_jaxpr to audit for effects before wrapping"],"tags":["jax","custom-jvp","effects","purity","jit"],"backgroundTag":"side-effects-in-traced-function","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}