{"record":{"id":"6584fce3d34f76cd","repo":"jax-ml/jax","slug":"no-jvp-defined-for-custom-jvp-function-self-f-n","errorCode":null,"errorMessage":"No JVP defined for custom_jvp function {self.f.__name__} using defjvp.","messagePattern":"No JVP defined for custom_jvp function (.+?) using defjvp\\.","errorType":"exception","errorClass":"AttributeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1198,"sourceCode":"\n  def defjvps(self, *jvps):\n    if self.static_argnums:\n      raise TypeError(\"Can't use ``defjvps`` with ``nondiff_argnums``.\")\n    def jvp(primals, tangents):\n      primal_out = self(*primals)\n      zeros = tree_map(ad_util.p2tz, primal_out)\n      all_tangents_out = [j(t, primal_out, *primals) if j else zeros\n                          for t, j in zip(tangents, jvps)]\n      sum_tangents = lambda _, x, *xs: reduce(ad.add_tangents, xs, x)\n      tangent_out = tree_map(sum_tangents, primal_out, *all_tangents_out)\n      return primal_out, tangent_out\n    self.defjvp(jvp)\n\n  def __call__(self, *args, **kwargs):\n    if not self.jvp_fun:\n      msg = (f\"No JVP defined for custom_jvp function {self.f.__name__} \"\n             \"using defjvp.\")\n      raise AttributeError(msg)\n\n    try:\n      args = resolve_kwargs(self.f, args, kwargs)\n    except TypeError as e:\n      raise TypeError(\n          \"The input arguments to the custom_jvp-decorated function \"\n          f\"{self.f.__name__} could not be resolved to positional-only \"\n          f\"arguments. Binding failed with the error:\\n{e}\") from e\n    if any(isinstance(args[i], core.Tracer) for i in self.static_argnums):\n      raise UnexpectedTracerError(\"custom_jvp inputs marked with nondiff_argnums \"\n                                  \"must be static, not Tracers\")\n    if all(is_hashable(args[i]) for i in self.static_argnums):\n      traced = api.jit(self.f, static_argnums=(*self.static_argnums,)).trace(*args)\n    else:\n      # jit requires hashable static_argnums values, but classic custom_jvp\n      # accepted unhashable nondiff_argnums values, so close over them instead\n      which_static = [i in self.static_argnums for i in range(len(args))]\n      dyn_args, static_args = partition_list(which_static, args)","sourceCodeStart":1180,"sourceCodeEnd":1216,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1180-L1216","documentation":"A @jax.custom_jvp-decorated function was called before its JVP rule was registered via defjvp (or @f.defjvp decorator). Like the custom_vjp analogue, the wrapper's __call__ raises AttributeError when self.jvp_fun is None.","triggerScenarios":"Calling f(x) where f is @jax.custom_jvp-decorated but f.defjvp(...) was never executed — e.g. registration is in a lazily imported module or was deleted during refactoring.","commonSituations":"Notebook workflows where cell defining the rule wasn't run; library code where registration happens in __init__py that got skipped; splitting definition and rule across modules with import cycles.","solutions":["Register the rule immediately after decoration: f.defjvp(f_jvp) at module scope","Ensure the module containing the registration is imported before use","Run cells in order in notebooks / verify registration in a smoke test"],"exampleFix":"// before\n@jax.custom_jvp\ndef f(x):\n  return x * 2\ny = f(3.)  # AttributeError\n// after\n@jax.custom_jvp\ndef f(x):\n  return x * 2\n@f.defjvp\ndef f_jvp(primals, tangents):\n  (x,), (xd,) = primals, tangents\n  return x * 2, 2 * xd\ny = f(3.)","handlingStrategy":"validation","validationCode":"assert getattr(f, 'jvp_fun', None) is not None, 'call f.defjvp(f_jvp) before invoking f'","typeGuard":null,"tryCatchPattern":"try:\n    f(x)\nexcept AttributeError as e:\n    if 'No JVP defined' in str(e):\n        f.defjvp(f_jvp); f(x)","preventionTips":["Register defjvp immediately after decoration at module scope","In notebooks, ensure the rule cell ran before calling"],"tags":["jax","custom-jvp","missing-rule","autodiff"],"backgroundTag":"jax-custom-jvp-missing-defjvp","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}