{"record":{"id":"4aff2f12960b3128","repo":"jax-ml/jax","slug":"no-jvp-defined-for-custom-jvp-function-primal-nam","errorCode":null,"errorMessage":"No JVP defined for custom_jvp function {primal_name} using defjvp.","messagePattern":"No JVP defined for custom_jvp function (.+?) using defjvp\\.","errorType":"exception","errorClass":"AttributeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":263,"sourceCode":"    def jvp(primals, tangents):\n      primal_out = self(*primals)\n      zeros = _zeros_like_pytree(primal_out)\n      all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros\n                          for t, jvp in zip(tangents, jvps)]\n      tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out)\n      return primal_out, tangent_out\n\n    self.defjvp(jvp)\n\n  @partial(traceback_util.api_boundary,\n           repro_api_name=\"jax.custom_jvp.__call__\")\n  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:\n    debug = debug_info(\"custom_jvp fun\", self.fun, args, kwargs,\n                       static_argnums=self.nondiff_argnums)\n    primal_name = debug.func_name\n    if not self.jvp:\n      msg = f\"No JVP defined for custom_jvp function {primal_name} using defjvp.\"\n      raise AttributeError(msg)\n\n    try:\n      args = resolve_kwargs(self.fun, args, kwargs)\n    except TypeError as e:\n      raise TypeError(\n          \"The input arguments to the custom_jvp-decorated function \"\n          f\"{primal_name} could not be resolved to positional-only arguments. \"\n          f\"Binding failed with the error:\\n{e}\"\n      ) from e\n\n    if self.nondiff_argnums:\n      args = tuple(_stop_gradient(x) if i in self.nondiff_argnums else x\n                   for i, x in enumerate(args))\n      diff_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]\n      f_, dyn_args = argnums_partial(lu.wrap_init(self.fun, debug_info=debug),\n                                     diff_argnums, args,\n                                     require_static_args_hashable=False)\n      static_args = [args[i] for i in self.nondiff_argnums]","sourceCodeStart":245,"sourceCodeEnd":281,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L245-L281","documentation":"A @jax.custom_jvp function must have a JVP rule registered via <fun>.defjvp (or defjvps) before it can be called; the wrapper checks for the rule up front and raises AttributeError if none exists.","triggerScenarios":"Decorating a function with @jax.custom_jvp but never registering f.defjvp / f.defjvps, then calling it (even outside grad).","commonSituations":"Forgetting the second decoration step; defjvp assigned conditionally or in dead code; merge conflicts dropping the rule block.","solutions":["Register a JVP rule with @f.defjvp immediately after the decorated function","If you don't need custom derivatives, remove the @jax.custom_jvp decorator","Ensure the defjvp block is at module level and executes at import"],"exampleFix":"# before\n@jax.custom_jvp\ndef f(x): return jnp.sin(x)\n# after\n@jax.custom_jvp\ndef f(x): return jnp.sin(x)\n@f.defjvp\ndef f_jvp(p, t):\n    (x,), (tx,) = p, t\n    return f(x), jnp.cos(x) * tx","handlingStrategy":"validation","validationCode":"assert getattr(f, 'jvp', None) is not None, 'register f.defjvp before use'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pair @jax.custom_jvp with a defjvp block in the same file","Run one grad call in tests to exercise the rule"],"tags":["jax","custom-jvp","autodiff","missing-rule"],"backgroundTag":"missing-callback-registration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}