{"record":{"id":"deefbe2453c63fd0","repo":"jax-ml/jax","slug":"vjp-from-jvp-is-a-pair-of-rules-not-a-single-ru","errorCode":null,"errorMessage":"`vjp_from_jvp` is a pair of rules, not a single rule; unpack it in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`","messagePattern":"`vjp_from_jvp` is a pair of rules, not a single rule; unpack it in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":671,"sourceCode":"  zero = lambda x: isinstance(x, ad_util.Zero)\n  out_ct = tree_map(ad_util.instantiate, out_ct, is_leaf=zero)\n  dummies = tree_map(lambda a: ad_util.zeros_like_aval(a.to_tangent_aval()),\n                     self.in_avals)\n  return api.linear_transpose(tangent_map, *dummies)(out_ct)\n\nclass _LinearizeFromJVP(NamedTuple):\n  lin: Callable\n  linearized: Callable\n  def __call__(self, *args, **kwargs):\n    raise TypeError(\n        \"`linearize_from_jvp` is a pair of rules, not a single rule; unpack \"\n        \"it in the class body: `lin, linearized = linearize_from_jvp`\")\n\nclass _VJPFromJVP(NamedTuple):\n  vjp_fwd: Callable\n  vjp_bwd_retval: Callable\n  def __call__(self, *args, **kwargs):\n    raise TypeError(\n        \"`vjp_from_jvp` is a pair of rules, not a single rule; unpack it in \"\n        \"the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`\")\n\nclass _VJPFromLin(NamedTuple):\n  vjp_fwd: Callable\n  vjp_bwd_retval: Callable\n  def __call__(self, *args, **kwargs):\n    raise TypeError(\n        \"`vjp_from_lin` is a pair of rules, not a single rule; unpack it in \"\n        \"the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_lin`\")\n\nlinearize_from_jvp = _LinearizeFromJVP(_lin_from_jvp, _linearized_from_jvp)\nvjp_from_jvp = _VJPFromJVP(_vjp_fwd_from_jvp, _transpose_jvp)\nvjp_from_lin = _VJPFromLin(_vjp_fwd_from_lin, _transpose_linearized)\n\n\nclass CustomVJPTraced(HiPrim):\n  \"\"\"Applications take ``(consts, fwd_consts, *args)``.","sourceCodeStart":653,"sourceCodeEnd":689,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L653-L689","documentation":"vjp_from_jvp is a NamedTuple pair (vjp_fwd, vjp_bwd_retval); it is not itself callable. Invoking it directly raises this TypeError pointing to the required class-body unpacking.","triggerScenarios":"Writing `vjp_fwd = vjp_from_jvp` or calling `vjp_from_jvp(res, ct)` instead of `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`.","commonSituations":"Migrating from classic custom_vjp where a single bwd function is registered.","solutions":["Unpack in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`"],"exampleFix":"# before\nclass P(hijax.HiPrim):\n  vjp_fwd = vjp_from_jvp\n# after\nclass P(hijax.HiPrim):\n  vjp_fwd, vjp_bwd_retval = vjp_from_jvp","handlingStrategy":"validation","validationCode":"from jax._src import hijax\nfor attr in ('vjp_fwd', 'vjp_bwd_retval'):\n    assert not isinstance(type(p).__dict__.get(attr), hijax._VJPFromJVP), attr + ' not unpacked'","typeGuard":"def vjp_rules_valid(p) -> bool:\n    from jax._src import hijax\n    d = type(p).__dict__\n    return not isinstance(d.get('vjp_fwd'), hijax._VJPFromJVP) and \\\n           not isinstance(d.get('vjp_bwd_retval'), hijax._VJPFromJVP) and \\\n           not isinstance(d.get('vjp_fwd'), hijax._VJPFromLin)","tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept TypeError as e:\n    if 'pair of rules' in str(e):\n        raise RuntimeError('unpack: vjp_fwd, vjp_bwd_retval = vjp_from_jvp') from e\n    raise","preventionTips":["Unpack pair derivations in the class body","Treat any module-level NamedTuple rule as a pair to unpack"],"tags":["jax","api-misuse","unpacking"],"backgroundTag":"rule-pair-not-unpacked","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}