{"record":{"id":"1c50c9ab010cb493","repo":"jax-ml/jax","slug":"type-prim-name-returned-structured-residua","errorCode":null,"errorMessage":"{type(_prim).__name__} returned structured residuals from `vjp_fwd`, which requires overriding `vjp_bwd(res, sres, outgrad, *arg_accums)`","messagePattern":"(.+?) returned structured residuals from `vjp_fwd`, which requires overriding `vjp_bwd\\(res, sres, outgrad, \\*arg_accums\\)`","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":426,"sourceCode":"# A `lin` or `vjp_fwd` rule may return (ans, res), (ans, res, nzs_out), or\n# (ans, res, nzs_out, sres). When it returns structured residuals, the paired\n# backward rule receives them explicitly: `linearized(res, sres, *tangents)`,\n# and `vjp_bwd(res, sres, outgrad, *arg_accums)` (which must be overridden).\ndef _call_hi_primitive_linearize(is_vjp, nz_in_flat, *args_flat, _prim):\n  args = tree_unflatten(_prim.in_tree, args_flat)\n  nzs_in = tree_unflatten(_prim.in_tree, nz_in_flat)\n  if is_vjp:\n    ans, residuals, *rest = _prim.vjp_fwd(nzs_in, *args)\n    linearized = partial(fake_linear_op, _prim, nz_in_flat)\n  else:\n    ans, residuals, *rest = _prim.lin(nzs_in, *args)\n    linearized = partial(flatten_user_linearized, _prim)\n  ans_flat = tree_leaves_checked(_prim.out_tree, ans)\n  nzs_out = rest[0] if rest else True\n  sres = rest[1] if len(rest) > 1 else None\n  if (sres is not None and is_vjp and\n      type(_prim).vjp_bwd is HiPrim.vjp_bwd):\n    raise TypeError(\n        f\"{type(_prim).__name__} returned structured residuals from `vjp_fwd`, \"\n        \"which requires overriding `vjp_bwd(res, sres, outgrad, *arg_accums)`\")\n  nzs_out_flat = broadcast_prefix(nzs_out, ans)\n  linearized = partial(linearized, nzs_out_flat) if is_vjp else linearized\n  return ans_flat, nzs_out_flat, residuals, sres, linearized\nad.primitive_linearizations[call_hi_primitive_p] = _call_hi_primitive_linearize\n\ndef fake_linear_op(prim, nz_in_flat, nz_out_flat, rs, sres, *tangents):\n  rs = rs if sres is None else (rs, sres)  # unpacked in the transpose rule\n  residuals_flat, residuals_tree = tree_flatten(rs)\n  assert nz_in_flat == [not isinstance(t, ad_util.Zero) for t in tangents]\n  nz_tangents = tree_leaves(tangents)\n  out_nz = call_hi_primitive_linearized_p.bind(\n      *residuals_flat, *nz_tangents, residuals_tree=residuals_tree, _prim=prim,\n      nz_in_flat=tuple(nz_in_flat), nz_out_flat=tuple(nz_out_flat),\n      has_sres=sres is not None)\n  out_nz_iter = iter(out_nz)\n  out = [next(out_nz_iter) if nz else ad_util.Zero(a.to_tangent_aval())","sourceCodeStart":408,"sourceCodeEnd":444,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L408-L444","documentation":"A HiPrim's vjp_fwd returned structured residuals (a second residual tree beyond flat residuals), which is only supported if the subclass overrides the accumulator-style `vjp_bwd(res, sres, outgrad, *arg_accums)` signature; the default HiPrim.vjp_bwd does not accept sres.","triggerScenarios":"A subclass vjp_fwd returns more than two values (out, residuals, sres...) and the class still uses the inherited vjp_bwd, during linearize staging of a vjp.","commonSituations":"Upgrading a custom primitive to keep structured residual state (e.g. per-layer caches) without updating the backward rule signature.","solutions":["Override `def vjp_bwd(self, res, sres, outgrad, *arg_accums)` on the subclass","Or stop returning structured residuals from vjp_fwd (return only flat residuals)","Or use vjp_bwd_retval with the matching sres-aware flatten path"],"exampleFix":"class MyPrim(hijax.HiPrim):\n  def vjp_fwd(self, *args):\n    out, res, sres = ...\n    return out, res, sres\n  # after: add sres-aware backward\n  def vjp_bwd(self, res, sres, outgrad, *accums):\n    ...","handlingStrategy":"type-guard","validationCode":"from jax._src.hijax import HiPrim\nassert not (returns_sres and type(prim).vjp_bwd is HiPrim.vjp_bwd), \\\n    'structured residuals require overriding vjp_bwd(res, sres, outgrad, *accums)'","typeGuard":"def sres_supported(p) -> bool:\n    from jax._src.hijax import HiPrim\n    return type(p).vjp_bwd is not HiPrim.vjp_bwd","tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept TypeError as e:\n    if 'structured residuals' in str(e):\n        raise RuntimeError('implement vjp_bwd with sres or drop sres from vjp_fwd') from e\n    raise","preventionTips":["When vjp_fwd returns >2 values, always add the sres-aware vjp_bwd","Document the residual layout of each custom primitive"],"tags":["jax","autodiff","vjp","api-contract"],"backgroundTag":"autodiff-rule-signature-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}