{"record":{"id":"b6923a03d80f1179","repo":"jax-ml/jax","slug":"type-prim-name-vjp-bwd-should-return-none","errorCode":null,"errorMessage":"{type(_prim).__name__}.vjp_bwd should return None or a dict of backward-pass log entries, got {type(log).__name__}","messagePattern":"(.+?)\\.vjp_bwd should return None or a dict of backward-pass log entries, got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":492,"sourceCode":"  accums_flat_ = iter(accums_flat)\n  accums_flat = [next(accums_flat_) if nz else ad.NullAccum(aval.to_ct_aval())\n                 for aval, nz in zip(_prim.in_avals_flat, nz_in_flat)]\n  assert next(accums_flat_, None) is None\n  accums = tree_unflatten(_prim.in_tree, accums_flat)\n  cts_flat_iter = iter(cts_flat_)\n  cts_flat = [next(cts_flat_iter) if nz else ad_util.Zero(a.to_ct_aval())\n              for a, nz in zip(_prim.out_avals_flat, nz_out_flat)]\n  assert next(cts_flat_iter, sentinel := object()) is sentinel\n  cts = tree_unflatten(_prim.out_tree, cts_flat)\n  # A vjp_bwd rule may return a dict of pytrees to log out of the backward\n  # pass (see VJP.with_logs), or None (the usual case) to log nothing.\n  if has_sres:\n    residuals, sres = residuals\n    log = _prim.vjp_bwd(residuals, sres, cts, *accums)\n  else:\n    log = _prim.vjp_bwd(residuals, cts, *accums)\n  if log is not None and type(log) is not dict:\n    raise TypeError(\n        f\"{type(_prim).__name__}.vjp_bwd should return None or a dict of \"\n        f\"backward-pass log entries, got {type(log).__name__}\")\n  return log\nad.fancy_transposes[call_hi_primitive_linearized_p] = _call_hi_primitive_linearized_transpose\n\ndef _call_hi_primitive_linearized_prettyprint(eqn, context, settings):\n  params = dict(eqn.params, _prim=eqn.params['_prim'].__class__.__name__,\n                residuals_tree='...')\n  if not params['has_sres']:\n    del params['has_sres']\n  return core._pp_eqn(eqn.replace(params=params), context, settings)\ncore.pp_eqn_rules[call_hi_primitive_linearized_p] = _call_hi_primitive_linearized_prettyprint\n\ndef _call_hi_primitive_jvp(primals, tangents, *, _prim):\n  primals = tree_unflatten(_prim.in_tree, primals)\n  tangents = tree_unflatten(_prim.in_tree, tangents)\n  out_primals, out_tangents = _prim.jvp(primals, tangents)\n  out_primals_flat = tree_leaves_checked(_prim.out_tree, out_primals)","sourceCodeStart":474,"sourceCodeEnd":510,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L474-L510","documentation":"The transpose of the linearized HiPrim primitive ran vjp_bwd, whose return value must be None or a dict of backward-pass log entries; anything else is a TypeError. The actual cotangents flow through accumulators, so the return is only for logging.","triggerScenarios":"A custom vjp_bwd implementation returns the input cotangents (a tuple/list/array) instead of None or a logs dict.","commonSituations":"Porting a classic custom_vjp bwd rule that returns in_cts into the hijax accumulator API without dropping the return value.","solutions":["Change vjp_bwd to return None (or a dict of logs) and deliver cotangents via the arg_accums accumulators","If you want the classic return-based API, implement vjp_bwd_retval instead","Return logs only as {key: value} entries for backward-pass logging"],"exampleFix":"# before\ndef vjp_bwd(self, res, outgrad, *accums):\n    return cts  # wrong\n# after\ndef vjp_bwd(self, res, outgrad, *accums):\n    ...\n    return None  # or {'loss_scale': ls}","handlingStrategy":"try-catch","validationCode":null,"typeGuard":"def vjp_bwd_returns_valid(p) -> bool:\n    out = p.vjp_bwd(dummy_res, dummy_ct, *dummy_accums)\n    return out is None or type(out) is dict","tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept TypeError as e:\n    if 'vjp_bwd should return' in str(e):\n        raise RuntimeError('vjp_bwd must return None or a logs dict; cots go via accumulators') from e\n    raise","preventionTips":["Return None from vjp_bwd unless logging","Use vjp_bwd_retval for return-based cotangents"],"tags":["jax","autodiff","vjp","api-contract"],"backgroundTag":"autodiff-rule-return-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}