{"record":{"id":"10129ab685aa8a60","repo":"jax-ml/jax","slug":"vjp-too-many-args-name-len-extra-args-1","errorCode":null,"errorMessage":"{_vjp_too_many_args(name, len(extra_args) + 1)}","messagePattern":"\\{_vjp_too_many_args\\(name, len\\(extra_args\\) \\+ 1\\)\\}","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":2028,"sourceCode":"  pass\n\n\n@dataclasses.dataclass(slots=True, weakref_slot=True)\nclass VJP:\n  fun: Callable  # partial(_vjp3_callable, ...)\n  in_tree: PyTreeDef\n  out_tree: PyTreeDef\n  args_res: list[Any]\n  opaque_residuals: list[Any]\n  structured_residuals: list[Any]\n  want_logs: bool = False\n  jaxpr = property(lambda self: self.fun.args[2])\n  out_nzs = property(lambda self: tuple(not z for z in self.fun.args[1]))\n\n  def __call__(self, out_ct, *extra_args):\n    if extra_args:\n      name, *_ = self.jaxpr.debug_info.func_src_info.split(' ')\n      raise TypeError(_vjp_too_many_args(name, len(extra_args) + 1))\n    return self.fun(self.in_tree, self.out_tree, self.args_res,\n                    self.opaque_residuals, self.structured_residuals,\n                    self.want_logs)(out_ct)\n\n  # Like __call__, but returns a pair (arg_cts, logs), where logs is a dict\n  # merging (with clobber semantics, in backward execution order) the dicts\n  # logged by transpose/vjp_bwd rules. Plain __call__ drops the logs.\n  with_logs = property(lambda self: self.replace(want_logs=True))\n\n  def with_refs(self, *maybe_ct_refs):\n    return self.fun(self.in_tree, self.out_tree, self.args_res,\n                    self.opaque_residuals, self.structured_residuals,\n                    self.want_logs, *maybe_ct_refs)\n\n  replace = dataclasses.replace\n\n  # Only safe to put these in cache keys if residuals aren't mutated. Beware!\n  __hash__ = object.__hash__","sourceCodeStart":2010,"sourceCodeEnd":2046,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L2010-L2046","documentation":"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.","triggerScenarios":"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.","commonSituations":"Confusion between jax.grad-style APIs and vjp pullbacks; unpacking outputs into separate variables and then passing them as separate arguments to the pullback.","solutions":["Wrap multiple cotangents in a single tuple: vjp_fn((ct1, ct2))","Pass a pytree with the same structure as the forward output as the single argument"],"exampleFix":"// before\nin_cts = vjp_fn(ct1, ct2)\n// after\nin_cts = vjp_fn((ct1, ct2))","handlingStrategy":"validation","validationCode":"out, vjp_fn = jax.vjp(f, *args)\n# vjp_fn takes ONE argument; wrap multiple cotangents in a tuple\ncall_vjp = lambda *cts: vjp_fn(cts if len(cts) > 1 else cts[0])","typeGuard":null,"tryCatchPattern":"try:\n    in_cts = vjp_fn(*cts)\nexcept TypeError:\n    in_cts = vjp_fn(cts)","preventionTips":["Remember the pullback signature: one cotangent pytree, not *args","Pass the primal output's structure directly as the cotangent template"],"tags":["jax","vjp","arity","argument-count"],"backgroundTag":"wrong-argument-count","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}