{"record":{"id":"9688aa4f186eca31","repo":"jax-ml/jax","slug":"vjp-from-lin-is-a-pair-of-rules-not-a-single-ru","errorCode":null,"errorMessage":"`vjp_from_lin` is a pair of rules, not a single rule; unpack it in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_lin`","messagePattern":"`vjp_from_lin` is a pair of rules, not a single rule; unpack it in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_lin`","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":679,"sourceCode":"  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)``.\n\n  The two leading arguments are synthetic, and both get zero cotangents:\n  ``consts`` is the primal function's closure environment promoted to an\n  argument (see ``Traced.with_consts_as_arg``), and ``fwd_consts`` is extra\n  inputs consumed only by the fwd rule and ignored by the primal (ordinarily\n  ``()``; the ``remat`` rule uses it to pass replay residuals to the helper\n  primitive it builds). Any value a rule needs beyond the primal arguments\n  must arrive through these slots as an explicit input, never by closure: the","sourceCodeStart":661,"sourceCodeEnd":697,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L661-L697","documentation":"vjp_from_lin is a NamedTuple pair (vjp_fwd, vjp_bwd_retval) derived from linearize rules; it is not directly callable, and calling it raises this TypeError instructing the class-body unpack.","triggerScenarios":"Writing `vjp_fwd = vjp_from_lin` or calling `vjp_from_lin(...)` directly rather than unpacking `vjp_fwd, vjp_bwd_retval = vjp_from_lin`.","commonSituations":"Copy-paste from docs where only one member of the pair was used; assuming tuple member access like vjp_from_lin.vjp_fwd is fine but direct calls are not.","solutions":["Unpack in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_lin` (requires lin/linearized to be defined, not derived from jvp)"],"exampleFix":"# before\nclass P(hijax.HiPrim):\n  vjp_bwd_retval = vjp_from_lin\n# after\nclass P(hijax.HiPrim):\n  lin, _ = linearize_from_jvp\n  vjp_fwd, vjp_bwd_retval = vjp_from_lin","handlingStrategy":"validation","validationCode":"from jax._src import hijax\nassert not isinstance(type(p).__dict__.get('vjp_fwd'), hijax._VJPFromLin)","typeGuard":"def vjp_from_lin_valid(p) -> bool:\n    from jax._src import hijax\n    return not isinstance(type(p).__dict__.get('vjp_fwd'), hijax._VJPFromLin) and \\\n           not isinstance(type(p).__dict__.get('vjp_bwd_retval'), 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_lin') from e\n    raise","preventionTips":["Unpack: vjp_fwd, vjp_bwd_retval = vjp_from_lin","Ensure lin/linearized are real rules, not derived"],"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"}