{"record":{"id":"94206eaf69cb153a","repo":"jax-ml/jax","slug":"custom-vjp-bwd-rule-self-bwd-must-produce-a-tupl-94206e","errorCode":null,"errorMessage":"Custom VJP bwd rule {self.bwd} must produce a tuple of length equal to the primal args tuple, but got length {len(in_cts)}","messagePattern":"Custom VJP bwd rule (.+?) must produce a tuple of length equal to the primal args tuple, but got length (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":801,"sourceCode":"        raise TypeError(\n            f\"Custom VJP bwd rule {self.bwd} was registered with \"\n            \"defvjp_with_logs and so must produce a pair (in_cts, logs), \"\n            f\"but got {in_cts}.\")\n      in_cts, logs = in_cts\n      if logs is not None and type(logs) is not dict:\n        raise TypeError(\n            f\"Custom VJP bwd rule {self.bwd} was registered with \"\n            \"defvjp_with_logs, and so the second element of the pair it \"\n            \"returns must be None or a dict of backward-pass log entries, \"\n            f\"but got {type(logs).__name__}.\")\n    if isinstance(in_cts, list):\n      in_cts = tuple(in_cts)\n    if not isinstance(in_cts, tuple):\n      raise TypeError(f\"Custom VJP bwd rule {self.bwd} must produce a tuple \"\n                      f\"but got {type(in_cts)}.\")\n    in_cts = (None, None, *in_cts)  # zero cts for the consts and fwd_consts args\n    if len(in_cts) != len(self.in_tree.children()) - len(self.static_argnums):\n      raise ValueError(f\"Custom VJP bwd rule {self.bwd} must produce a tuple \"\n                       \"of length equal to the primal args tuple, but got \"\n                       f\"length {len(in_cts)}\")\n    in_cts = broadcast_prefix(in_cts, in_avals_, is_leaf=lambda x: x is None)\n    in_cts = tree_unflatten(self.in_tree, map(_replace_none, self.in_avals_flat, in_cts))\n    tree_map_with_path(partial(_vjp_bwd_aval_mismatch_err, self.traced._fun_sourceinfo),\n                               self.in_avals[2:], in_cts[2:])\n    if self.symbolic_zeros:\n      in_cts = tree_map(ad_util.replace_rule_output_symbolic_zeros, in_cts)\n    return (in_cts, logs) if self.with_logs else in_cts\n\n  def jvp(self, primals, tangents):\n    if self.symbolic_zeros: ad.raise_custom_vjp_error_on_jvp()\n    zero = lambda x: isinstance(x, ad_util.Zero)\n    nzs_in = tuple(tree_map(lambda t: not isinstance(t, ad_util.Zero), t,\n                            is_leaf=zero) for t in tangents)\n    tangents = tree_map(ad_util.instantiate, tangents, is_leaf=zero)\n    if self.opt_remat:\n      fwd_traced = api.jit(partial(self.vjp_fwd, nzs_in)).trace(*primals)","sourceCodeStart":783,"sourceCodeEnd":819,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L783-L819","documentation":"The custom VJP backward rule returned a tuple whose length does not match the number of primal (non-static) arguments after two leading None entries for consts/fwd_consts are inserted. ValueError raised during length validation.","triggerScenarios":"bwd returns more or fewer cotangents than the function has non-static arguments (e.g. two cts for a three-arg function, or forgetting that some args are static).","commonSituations":"Adding/removing function arguments after writing the bwd rule; static_argnums in the custom_vjp registration shrinking the expected count.","solutions":["Return exactly one cotangent (possibly None) per non-static primal argument","Re-count args after signature changes, including defaults and static args","Use jax.tree.map over the input tree to build the return"],"exampleFix":"# before (3 dynamic args)\ndef bwd(res, ct):\n    return (g1,)\n# after\ndef bwd(res, ct):\n    return (g1, g2, g3)","handlingStrategy":"validation","validationCode":"n_expected = len(prim.in_tree.children()) - len(prim.static_argnums) - 2\ncts = bwd(*static_args, res, out_ct)[0] if with_logs else bwd(...)\nassert len(cts) == n_expected, f'expected {n_expected} cotangents, got {len(cts)}'","typeGuard":"def cts_arity_ok(cts, n_dynamic_args) -> bool:\n    return len(cts) == n_dynamic_args","tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept ValueError as e:\n    if 'length equal to the primal args tuple' in str(e):\n        raise RuntimeError('bwd must return one ct per non-static primal arg') from e\n    raise","preventionTips":["Update bwd whenever the primal signature changes","Track static_argnums carefully; they reduce the expected ct count","Return None placeholders for args that receive no cotangent"],"tags":["jax","custom-vjp","arity-mismatch"],"backgroundTag":"custom-vjp-return-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}