jax-ml/jax · error · AttributeError

No JVP defined for custom_jvp function {self.f.__name__} usi

Error message

No JVP defined for custom_jvp function {self.f.__name__} using defjvp.

What it means

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.

Source

Thrown at jax/_src/hijax.py:1198

  def defjvps(self, *jvps):
    if self.static_argnums:
      raise TypeError("Can't use ``defjvps`` with ``nondiff_argnums``.")
    def jvp(primals, tangents):
      primal_out = self(*primals)
      zeros = tree_map(ad_util.p2tz, primal_out)
      all_tangents_out = [j(t, primal_out, *primals) if j else zeros
                          for t, j in zip(tangents, jvps)]
      sum_tangents = lambda _, x, *xs: reduce(ad.add_tangents, xs, x)
      tangent_out = tree_map(sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out
    self.defjvp(jvp)

  def __call__(self, *args, **kwargs):
    if not self.jvp_fun:
      msg = (f"No JVP defined for custom_jvp function {self.f.__name__} "
             "using defjvp.")
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.f, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "The input arguments to the custom_jvp-decorated function "
          f"{self.f.__name__} could not be resolved to positional-only "
          f"arguments. Binding failed with the error:\n{e}") from e
    if any(isinstance(args[i], core.Tracer) for i in self.static_argnums):
      raise UnexpectedTracerError("custom_jvp inputs marked with nondiff_argnums "
                                  "must be static, not Tracers")
    if all(is_hashable(args[i]) for i in self.static_argnums):
      traced = api.jit(self.f, static_argnums=(*self.static_argnums,)).trace(*args)
    else:
      # jit requires hashable static_argnums values, but classic custom_jvp
      # accepted unhashable nondiff_argnums values, so close over them instead
      which_static = [i in self.static_argnums for i in range(len(args))]
      dyn_args, static_args = partition_list(which_static, args)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Register the rule immediately after decoration: f.defjvp(f_jvp) at module scope
2. Ensure the module containing the registration is imported before use
3. Run cells in order in notebooks / verify registration in a smoke test

Example fix

// before
@jax.custom_jvp
def f(x):
  return x * 2
y = f(3.)  # AttributeError
// after
@jax.custom_jvp
def f(x):
  return x * 2
@f.defjvp
def f_jvp(primals, tangents):
  (x,), (xd,) = primals, tangents
  return x * 2, 2 * xd
y = f(3.)

assert getattr(f, 'jvp_fun', None) is not None, 'call f.defjvp(f_jvp) before invoking f'

try:
    f(x)
except AttributeError as e:
    if 'No JVP defined' in str(e):
        f.defjvp(f_jvp); f(x)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• No JVP defined for custom_jvp function {primal_name} using d
• 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/6584fce3d34f76cd. Report an issue: GitHub.