{"record":{"id":"4a5241be5db1eca7","repo":"jax-ml/jax","slug":"effects-not-supported-in-custom-jvp-disallowed-4a5241","errorCode":null,"errorMessage":"Effects not supported in `custom_jvp`: {disallowed}","messagePattern":"Effects not supported in `custom_jvp`: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":841,"sourceCode":"    nzs_in_flat = [True] * len(self.in_avals_flat)\n    nzs_out_flat = [True] * len(self.out_avals_flat)\n    tangents_flat = tree_leaves_checked(self.in_tree, tangents)\n    tangents_out_flat = fake_linear_op(self, nzs_in_flat, nzs_out_flat, residuals,\n                                       None, *tangents_flat)\n    tangents_out = tree_unflatten(self.out_tree, tangents_out_flat)\n    return primals_out, tangents_out\n\n  def batch_dim_rule(self, axis_data, in_dims):\n    _, primal_in_tree = tracing_registry.flatten(self.drop_fwd_consts(*self.in_avals))\n    in_dims_flat = primal_in_tree.flatten_up_to(self.drop_fwd_consts(*in_dims))\n    _, out_dims = batching.batch_jaxpr2(self.traced.jaxpr, axis_data, tuple(in_dims_flat))\n    return tree_unflatten(self.out_tree, out_dims)\n\n  def check(self, *_):\n    effs = self.traced.jaxpr.effects\n    disallowed = effects.custom_derivatives_allowed_effects.filter_not_in(effs)\n    if disallowed:\n      raise NotImplementedError(f'Effects not supported in `custom_jvp`: {disallowed}')\n\n  def remat(self, trace, *args):  # type: ignore\n    if self.opt_remat:\n      return self(*args), self\n    if not trace.custom_vjp_rules:\n      return self(*args), self  # see https://github.com/jax-ml/jax/pull/38914\n    if not self.static_argnums:\n      fwd, dyn_args = self.fwd, args\n    else:\n      which_static = [i in self.static_argnums for i in range(len(args))]\n      dyn_args, static_args = partition_list(which_static, args)\n      static_args = [x.val for x in static_args]\n      fwd = lambda *dyn_args: self.fwd(*merge_lists(which_static, list(dyn_args), static_args))\n    # custom_vjp_rules=False so that custom_vjp applications inside fwd hit\n    # the early return above rather than recursively tracing their fwds.\n    (out, _), rem_ = remat.remat_transform(trace.policy, fwd, *dyn_args,\n                                           custom_vjp_rules=False)\n    res = tuple(rem_.args[0])","sourceCodeStart":823,"sourceCodeEnd":859,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L823-L859","documentation":"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.","triggerScenarios":"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.","commonSituations":"Adding debug printing or state effects inside a custom-derivative function; upgrading JAX versions where effect tracking became stricter for custom_jvp/custom_vjp.","solutions":["Move effectful operations (prints, callbacks) outside the custom_jvp-decorated function","Remove or gate debug prints behind a non-traced code path","If you need effects with custom derivatives, restructure so effects occur in the caller, not inside the decorated function"],"exampleFix":"// before\n@jax.custom_jvp\ndef f(x):\n  jax.debug.print(\"x={}\", x)\n  return x * 2\n// after\n@jax.custom_jvp\ndef f(x):\n  return x * 2\n# print outside, or via a custom JVP rule that does not carry effects","handlingStrategy":"validation","validationCode":"# Keep effectful ops out of custom_jvp bodies; assert at registration time\nimport jax, jax.numpy as jnp\n@jax.custom_jvp\ndef f(x):\n    return x * 2\n# smoke-check the rule traces cleanly under grad before real use\njax.grad(lambda x: f(x).sum())(jnp.zeros(2))","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(loss_fn)(x)\nexcept NotImplementedError as e:\n    if 'Effects not supported in `custom_jvp`' in str(e):\n        # remove prints/callbacks from the decorated function and retry\n        ...","preventionTips":["Never put jax.debug.print or host callbacks inside custom_jvp/custom_vjp functions","Run a tiny grad smoke test after decorating","Keep custom-derivative functions pure numerical kernels"],"tags":["jax","custom-jvp","autodiff","effects"],"backgroundTag":"jax-custom-derivative-effects-not-supported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}