jax-ml/jax · error · NotImplementedError

Effects not supported in `custom_jvp`: {disallowed_effects}

Error message

Effects not supported in `custom_jvp`: {disallowed_effects}

What it means

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.

Source

Thrown at jax/_src/custom_derivatives.py:437

    out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents)
    out_primals, nz_out_tangents = split_list(out, [len(out_zeros)])
    nz_out_tangents_ = iter(nz_out_tangents)
    out_tangents = [SymbolicZero(core.typeof(p).to_tangent_aval())
                    if z else next(nz_out_tangents_)
                    for p, z in zip(out_primals, out_zeros)]
    assert next(nz_out_tangents_, None) is None
    return [*out_primals, *out_tangents]
  return lu.wrap_init(jvp, debug_info=jvp_jaxpr_fun.debug_info)

custom_jvp_call_p = CustomJVPCallPrimitive('custom_jvp_call')

def _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_fun,
                               num_consts, symbolic_zeros):
  # TODO(mattjj): could do more checking here...
  del in_avals, jvp_jaxpr_fun, num_consts
  disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(call_jaxpr.effects)
  if disallowed_effects:
    raise NotImplementedError(
        f'Effects not supported in `custom_jvp`: {disallowed_effects}')
  return call_jaxpr.out_avals, core.positional_effects(call_jaxpr)
core.custom_typechecks[custom_jvp_call_p] = _custom_jvp_call_typecheck

def _custom_jvp_vjp_call_lowering(ctx: mlir.LoweringRuleContext, *args,
                                  call_jaxpr: core.Jaxpr, **_):
  consts = mlir.ir_consts(
      call_jaxpr.consts, [v.aval for v in call_jaxpr.constvars])
  out, tokens = mlir.jaxpr_subcomp(ctx.module_context, call_jaxpr,
                                   ctx.name_stack, ctx.tokens_in, consts,
                                   *args, dim_var_values=ctx.dim_var_values,
                                   const_lowering=ctx.const_lowering,
                                   outer_traceback=ctx.traceback)
  ctx.set_tokens_out(tokens)
  return out
mlir.register_lowering(custom_jvp_call_p, _custom_jvp_vjp_call_lowering)

def _custom_jvp_call_transpose_fancy(params, jaxpr, args, ct, _):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move effectful operations (randomness, state reads/writes) outside the custom_jvp function and pass their results in as arguments
2. Use pure jax.random PRNG keys passed as arguments instead of stateful RNG
3. Remove debug callbacks/halts from the decorated body

Example fix

# before
@jax.custom_jvp
def f(x):
    key = next_key()  # effectful
    return x + jax.random.normal(key, x.shape)
# after
def make(x, key):
    return x + jax.random.normal(key, x.shape)
@jax.custom_jvp
def f(x, noise):
    return x + noise
# caller draws noise with an explicit key beforehand

jaxpr = jax.make_jaxpr(f)(*sample_args)
assert not jaxpr.effects, f'effects present: {jaxpr.effects}'

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Effects not supported in `custom_jvp`: {disallowed}
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Formatting arguments to checkify.check need to be PyTrees of
• Value of type {type(self)} is not convertible to float.
• Value of type {type(self)} is not convertible to complex.

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