jax-ml/jax · error · AttributeError

No JVP defined for custom_jvp function {primal_name} using d

Error message

No JVP defined for custom_jvp function {primal_name} using defjvp.

What it means

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.

Source

Thrown at jax/_src/custom_derivatives.py:263

    def jvp(primals, tangents):
      primal_out = self(*primals)
      zeros = _zeros_like_pytree(primal_out)
      all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros
                          for t, jvp in zip(tangents, jvps)]
      tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out

    self.defjvp(jvp)

  @partial(traceback_util.api_boundary,
           repro_api_name="jax.custom_jvp.__call__")
  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:
    debug = debug_info("custom_jvp fun", self.fun, args, kwargs,
                       static_argnums=self.nondiff_argnums)
    primal_name = debug.func_name
    if not self.jvp:
      msg = f"No JVP defined for custom_jvp function {primal_name} using defjvp."
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.fun, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "The input arguments to the custom_jvp-decorated function "
          f"{primal_name} could not be resolved to positional-only arguments. "
          f"Binding failed with the error:\n{e}"
      ) from e

    if self.nondiff_argnums:
      args = tuple(_stop_gradient(x) if i in self.nondiff_argnums else x
                   for i, x in enumerate(args))
      diff_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]
      f_, dyn_args = argnums_partial(lu.wrap_init(self.fun, debug_info=debug),
                                     diff_argnums, args,
                                     require_static_args_hashable=False)
      static_args = [args[i] for i in self.nondiff_argnums]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Register a JVP rule with @f.defjvp immediately after the decorated function
2. If you don't need custom derivatives, remove the @jax.custom_jvp decorator
3. Ensure the defjvp block is at module level and executes at import

Example fix

# before
@jax.custom_jvp
def f(x): return jnp.sin(x)
# after
@jax.custom_jvp
def f(x): return jnp.sin(x)
@f.defjvp
def f_jvp(p, t):
    (x,), (tx,) = p, t
    return f(x), jnp.cos(x) * tx

assert getattr(f, 'jvp', None) is not None, 'register f.defjvp before use'

When it happens

Trigger: Decorating a function with @jax.custom_jvp but never registering f.defjvp / f.defjvps, then calling it (even outside grad).

Common situations: Forgetting the second decoration step; defjvp assigned conditionally or in dead code; merge conflicts dropping the rule block.

Related errors

• No JVP defined for custom_jvp function {self.f.__name__} usi
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Can't use ``defjvps`` with ``nondiff_argnums``.
• 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/4aff2f12960b3128. Report an issue: GitHub.