jax-ml/jax · error · TypeError

{_vjp_too_many_args(name, len(extra_args) + 1)}

Error message

{_vjp_too_many_args(name, len(extra_args) + 1)}

What it means

The backward (VJP) function returned by jax.vjp takes exactly one argument: the cotangent pytree matching the function output. This TypeError fires when it is called with more than one positional argument, as if it were the forward function or a multi-arg pullback.

Source

Thrown at jax/_src/api.py:2028

  pass


@dataclasses.dataclass(slots=True, weakref_slot=True)
class VJP:
  fun: Callable  # partial(_vjp3_callable, ...)
  in_tree: PyTreeDef
  out_tree: PyTreeDef
  args_res: list[Any]
  opaque_residuals: list[Any]
  structured_residuals: list[Any]
  want_logs: bool = False
  jaxpr = property(lambda self: self.fun.args[2])
  out_nzs = property(lambda self: tuple(not z for z in self.fun.args[1]))

  def __call__(self, out_ct, *extra_args):
    if extra_args:
      name, *_ = self.jaxpr.debug_info.func_src_info.split(' ')
      raise TypeError(_vjp_too_many_args(name, len(extra_args) + 1))
    return self.fun(self.in_tree, self.out_tree, self.args_res,
                    self.opaque_residuals, self.structured_residuals,
                    self.want_logs)(out_ct)

  # Like __call__, but returns a pair (arg_cts, logs), where logs is a dict
  # merging (with clobber semantics, in backward execution order) the dicts
  # logged by transpose/vjp_bwd rules. Plain __call__ drops the logs.
  with_logs = property(lambda self: self.replace(want_logs=True))

  def with_refs(self, *maybe_ct_refs):
    return self.fun(self.in_tree, self.out_tree, self.args_res,
                    self.opaque_residuals, self.structured_residuals,
                    self.want_logs, *maybe_ct_refs)

  replace = dataclasses.replace

  # Only safe to put these in cache keys if residuals aren't mutated. Beware!
  __hash__ = object.__hash__

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap multiple cotangents in a single tuple: vjp_fn((ct1, ct2))
2. Pass a pytree with the same structure as the forward output as the single argument

Example fix

// before
in_cts = vjp_fn(ct1, ct2)
// after
in_cts = vjp_fn((ct1, ct2))

out, vjp_fn = jax.vjp(f, *args)
# vjp_fn takes ONE argument; wrap multiple cotangents in a tuple
call_vjp = lambda *cts: vjp_fn(cts if len(cts) > 1 else cts[0])

try:
    in_cts = vjp_fn(*cts)
except TypeError:
    in_cts = vjp_fn(cts)

When it happens

Trigger: Calling vjp_fn(ct1, ct2) instead of vjp_fn((ct1, ct2)) when the primal output is a 2-tuple; splatting a tuple of cotangents with * instead of passing it as one pytree.

Common situations: Confusion between jax.grad-style APIs and vjp pullbacks; unpacking outputs into separate variables and then passing them as separate arguments to the pullback.

Related errors

• reduce_axes argument to vjp is deprecated
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• But the tree structures differ:

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