{"record":{"id":"c49a97b4da58e90f","repo":"jax-ml/jax","slug":"custom-vjp-bwd-rule-self-bwd-was-registered-with","errorCode":null,"errorMessage":"Custom VJP bwd rule {self.bwd} was registered with defvjp_with_logs and so must produce a pair (in_cts, logs), but got {in_cts}.","messagePattern":"Custom VJP bwd rule (.+?) was registered with defvjp_with_logs and so must produce a pair \\(in_cts, logs\\), but got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":783,"sourceCode":"      out_nzs = tree_unflatten(self.out_tree, out_nzs_flat)\n      out = tree_unflatten(self.out_tree, out_flat)\n      return out, res, out_nzs\n    else:\n      return out, res\n\n  def vjp_bwd_retval(self, res, out_ct):\n    static_args = tuple(x.val for x in self.in_avals if isinstance(x, Static))\n    in_avals_ = tuple(x for x in self.in_avals if not isinstance(x, Static))\n    leaf = lambda x: isinstance(x, ad_util.Zero)\n    if self.symbolic_zeros:\n      out_ct = tree_map(ad_util.replace_internal_symbolic_zeros, out_ct, is_leaf=leaf)\n    else:\n      out_ct = tree_map(ad_util.instantiate, out_ct, is_leaf=leaf)\n    in_cts = self.bwd(*static_args, res, out_ct)\n    logs = None\n    if self.with_logs:\n      if not (isinstance(in_cts, (list, tuple)) and len(in_cts) == 2):\n        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 \"","sourceCodeStart":765,"sourceCodeEnd":801,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L765-L801","documentation":"A custom VJP backward rule registered with hijax's defvjp_with_logs must return a pair (in_cts, logs). The wrapper checked the bwd return value and found it was not a 2-element list/tuple.","triggerScenarios":"Calling grad/vjp over a function whose custom_vjp was registered via defvjp_with_logs, where the bwd function returns only in_cts (or some other shape).","commonSituations":"Upgrading existing custom_vjp code to defvjp_with_logs but forgetting to add the logs element to the bwd return.","solutions":["Return `(in_cts, None)` or `(in_cts, logs_dict)` from the bwd rule","If logging is not needed, register with plain defvjp instead of defvjp_with_logs"],"exampleFix":"# before\ndef bwd(res, ct):\n    return cts\n# after\ndef bwd(res, ct):\n    return cts, None","handlingStrategy":"try-catch","validationCode":"ret = bwd(*static_args, res, out_ct)\nassert isinstance(ret, (list, tuple)) and len(ret) == 2, 'defvjp_with_logs bwd must return (in_cts, logs)'","typeGuard":"def bwd_returns_pair(bwd, static_args, res, ct) -> bool:\n    r = bwd(*static_args, res, ct)\n    return isinstance(r, (list, tuple)) and len(r) == 2","tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept TypeError as e:\n    if 'must produce a pair' in str(e):\n        raise RuntimeError('bwd registered with logs must return (in_cts, logs)') from e\n    raise","preventionTips":["When using defvjp_with_logs, always return two elements","Prefer plain defvjp unless logs are needed"],"tags":["jax","custom-vjp","api-contract","logging"],"backgroundTag":"custom-vjp-return-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}