jax-ml/jax · error · NotImplementedError

Effects not supported in `custom_jvp`: {disallowed}

Error message

Effects not supported in `custom_jvp`: {disallowed}

What it means

Raised when a @jax.custom_jvp-decorated function's traced jaxpr contains effects (e.g. ordered effects like print or state) that are not allowed under custom derivatives. During JVP tracing, JAX filters the jaxpr's effects against custom_derivatives_allowed_effects and refuses to proceed if any remain, because autodiff rules cannot soundly replay side effects.

Source

Thrown at jax/_src/hijax.py:841

    nzs_in_flat = [True] * len(self.in_avals_flat)
    nzs_out_flat = [True] * len(self.out_avals_flat)
    tangents_flat = tree_leaves_checked(self.in_tree, tangents)
    tangents_out_flat = fake_linear_op(self, nzs_in_flat, nzs_out_flat, residuals,
                                       None, *tangents_flat)
    tangents_out = tree_unflatten(self.out_tree, tangents_out_flat)
    return primals_out, tangents_out

  def batch_dim_rule(self, axis_data, in_dims):
    _, primal_in_tree = tracing_registry.flatten(self.drop_fwd_consts(*self.in_avals))
    in_dims_flat = primal_in_tree.flatten_up_to(self.drop_fwd_consts(*in_dims))
    _, out_dims = batching.batch_jaxpr2(self.traced.jaxpr, axis_data, tuple(in_dims_flat))
    return tree_unflatten(self.out_tree, out_dims)

  def check(self, *_):
    effs = self.traced.jaxpr.effects
    disallowed = effects.custom_derivatives_allowed_effects.filter_not_in(effs)
    if disallowed:
      raise NotImplementedError(f'Effects not supported in `custom_jvp`: {disallowed}')

  def remat(self, trace, *args):  # type: ignore
    if self.opt_remat:
      return self(*args), self
    if not trace.custom_vjp_rules:
      return self(*args), self  # see https://github.com/jax-ml/jax/pull/38914
    if not self.static_argnums:
      fwd, dyn_args = self.fwd, args
    else:
      which_static = [i in self.static_argnums for i in range(len(args))]
      dyn_args, static_args = partition_list(which_static, args)
      static_args = [x.val for x in static_args]
      fwd = lambda *dyn_args: self.fwd(*merge_lists(which_static, list(dyn_args), static_args))
    # custom_vjp_rules=False so that custom_vjp applications inside fwd hit
    # the early return above rather than recursively tracing their fwds.
    (out, _), rem_ = remat.remat_transform(trace.policy, fwd, *dyn_args,
                                           custom_vjp_rules=False)
    res = tuple(rem_.args[0])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move effectful operations (prints, callbacks) outside the custom_jvp-decorated function
2. Remove or gate debug prints behind a non-traced code path
3. If you need effects with custom derivatives, restructure so effects occur in the caller, not inside the decorated function

Example fix

// before
@jax.custom_jvp
def f(x):
  jax.debug.print("x={}", x)
  return x * 2
// after
@jax.custom_jvp
def f(x):
  return x * 2
# print outside, or via a custom JVP rule that does not carry effects

# Keep effectful ops out of custom_jvp bodies; assert at registration time
import jax, jax.numpy as jnp
@jax.custom_jvp
def f(x):
    return x * 2
# smoke-check the rule traces cleanly under grad before real use
jax.grad(lambda x: f(x).sum())(jnp.zeros(2))

try:
    jax.grad(loss_fn)(x)
except NotImplementedError as e:
    if 'Effects not supported in `custom_jvp`' in str(e):
        # remove prints/callbacks from the decorated function and retry
        ...

When it happens

Trigger: Calling a @custom_jvp function (or differentiating through it with jax.grad/jax.jvp) whose body usesjax.debug.print, host callbacks, or other effectful primitives inside the primal computation that the custom JVP rule cannot account for.

Common situations: Adding debug printing or state effects inside a custom-derivative function; upgrading JAX versions where effect tracking became stricter for custom_jvp/custom_vjp.

Related errors

• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Can't use ``defjvps`` with ``nondiff_argnums``.
• No JVP defined for custom_jvp function {primal_name} using d
• Custom JVP rule {jvp_name} for function {primal_name} must p
• Custom JVP rule {jvp_name} for function {primal_name} must p

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