jax-ml/jax · error · AttributeError

No VJP defined for custom_vjp function {self.f.__name__} usi

Error message

No VJP defined for custom_vjp function {self.f.__name__} using defvjp.

What it means

A @jax.custom_vjp-decorated function was called before any VJP rule was registered with defvjp. The wrapper raises AttributeError on __call__ because neither self.fwd nor self.bwd exists yet.

Source

Thrown at jax/_src/hijax.py:944

    update_wrapper(self, f)
    self.f = f

  def defvjp(self, fwd, bwd, *, symbolic_zeros=False, optimize_remat=False):
    self.fwd = fwd
    self.bwd = bwd
    self.symz = symbolic_zeros
    self.opt_remat = optimize_remat

  def defvjp_with_logs(self, fwd, bwd, *, symbolic_zeros=False,
                       optimize_remat=False):
    self.defvjp(fwd, bwd, symbolic_zeros=symbolic_zeros,
                optimize_remat=optimize_remat)
    self.with_logs = True

  def __call__(self, *args, **kwargs):
    if not self.fwd or not self.bwd:
      msg = f"No VJP defined for custom_vjp function {self.f.__name__} using defvjp."
      raise AttributeError(msg)

    args = resolve_kwargs(self.f, args, kwargs)
    if any(isinstance(l, core.Tracer) for i in self.static_argnums
           for l in tree_leaves(args[i])):
      raise UnexpectedTracerError("custom_vjp 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_vjp
      # 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)
      f = dyn_args_fun(self.f, self.static_argnums,
                       tuple(map(WrapHashably, static_args)), len(args))
      traced = api.jit(f).trace(*dyn_args)
    args = tuple(Static(x) if i in self.static_argnums else x for i, x in enumerate(args))
    consts, traced = traced.with_consts_as_arg()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Register the rule before any call: f.defvjp(fwd, bwd)
2. If using the newer split API, ensure both defvjp_fwd and defvjp_bwd were invoked
3. Move registration to module top level right after the decorated definition so import order guarantees it

Example fix

// before
@jax.custom_vjp
def f(x):
  return x * 2
y = f(3.)   # AttributeError
// after
@jax.custom_vjp
def f(x):
  return x * 2
def f_fwd(x): return f(x), None
def f_bwd(_, ct): return (ct * 2,)
f.defvjp(f_fwd, f_bwd)
y = f(3.)

assert getattr(f, 'fwd', None) and getattr(f, 'bwd', None), 'call f.defvjp(fwd, bwd) before use'

try:
    f(x)
except AttributeError as e:
    if 'No VJP defined' in str(e):
        f.defvjp(f_fwd, f_bwd)
        f(x)

When it happens

Trigger: Calling f(x) on a function decorated with @jax.custom_vjp without ever calling f.defvjp(f_fwd, f_bwd) (or defvjp_fwd/defvjp_bwd) beforehand.

Common situations: Refactoring where defvvp registration moved below the first call site or into a module that is no longer imported; copy-paste from custom_jvp examples where registration is optional; defining the decorated function in a library where the user must register the rule themselves.

Related errors

• Pure callbacks do not support transpose. Please use `jax.cus
• No JVP defined for custom_jvp function {primal_name} using d
• at {keystr(path)}, got fwd output type {ty.str_short()} whic
• {result}the bwd rule attached to {primal_sourceinfo} produce
• No JVP defined for custom_jvp function {self.f.__name__} usi

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/c8dcf57ce70e319e. Report an issue: GitHub.