{"record":{"id":"78f681ae331bda38","repo":"jax-ml/jax","slug":"subclass-type-self-can-t-set-both-jvp-jvp-fr","errorCode":null,"errorMessage":"subclass {type(self)} can't set both `jvp = jvp_from_lin` and `lin, linearized = linearize_from_jvp`, since each would be defined in terms of the other","messagePattern":"subclass (.+?) can't set both `jvp = jvp_from_lin` and `lin, linearized = linearize_from_jvp`, since each would be defined in terms of the other","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":603,"sourceCode":"    return out_primals_flat, out_tangents_flat\n\n  dbg = debug_info('linearize_from_jvp', self.jvp, (primals, primals), {})\n  out_primals_flat, nzs_out_flat, consts, _, linearized = ad.linearize_from_jvp(\n      lu.wrap_init(jvp_flat, debug_info=dbg), True, nzs_in_flat,\n      False, False, primals_flat, {})\n  out_primals = tree_unflatten(self.out_tree, out_primals_flat)\n  nzs_out = tree_unflatten(self.out_tree, list(nzs_out_flat))\n  return out_primals, DerivedLinearization(consts, linearized), nzs_out\n\ndef _linearized_from_jvp(self, residuals, *tangents):\n  \"\"\"The `linearized` half of the `linearize_from_jvp` pair.\"\"\"\n  tangents_flat = self.in_tree.flatten_up_to(tangents)\n  out_tangents_flat = residuals.apply(residuals.consts, None, *tangents_flat)\n  return tree_unflatten(self.out_tree, out_tangents_flat)\n\ndef jvp_from_lin(self, primals, tangents):\n  if type(self).lin is _lin_from_jvp:\n    raise TypeError(\n        f\"subclass {type(self)} can't set both `jvp = jvp_from_lin` and \"\n        \"`lin, linearized = linearize_from_jvp`, since each would be defined \"\n        \"in terms of the other\")\n  tangents_flat = self.in_tree.flatten_up_to(tangents)\n  nzs_in = tree_unflatten(\n      self.in_tree, [not isinstance(t, ad_util.Zero) for t in tangents_flat])\n  out_primals, residuals, *rest = self.lin(nzs_in, *primals)\n  out_tangents = (self.linearized(residuals, *tangents) if len(rest) < 2 else\n                  self.linearized(residuals, rest[1], *tangents))\n  return out_primals, out_tangents\n\ndef _vjp_fwd_from_jvp(self, nzs_in, *primals):\n  \"\"\"The `vjp_fwd` half of the `vjp_from_jvp` pair.\"\"\"\n  return self(*primals), (primals, nzs_in)\n\ndef _transpose_jvp(self, res, out_ct):\n  \"\"\"The `vjp_bwd_retval` half of the `vjp_from_jvp` pair.\"\"\"\n  primals, nzs_in = res","sourceCodeStart":585,"sourceCodeEnd":621,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L585-L621","documentation":"jvp_from_lin detects that the class also set `lin, linearized = linearize_from_jvp`, i.e. both forward rules would be mutually defined in terms of each other, causing infinite recursion; it raises TypeError immediately.","triggerScenarios":"In a HiPrim class body setting both `jvp = jvp_from_lin` and `lin, linearized = linearize_from_jvp`.","commonSituations":"Copy-pasting all derivation decorators from the hijax docs without realizing these two are alternative directions (jvp-from-linearize vs linearize-from-jvp).","solutions":["Keep exactly one direction: either define jvp and use linearize_from_jvp, or define lin/linearized and use jvp_from_lin","Remove `jvp = jvp_from_lin` if you have real lin/linearized rules","Remove `lin, linearized = linearize_from_jvp` if you have a real jvp rule"],"exampleFix":"# before\nclass P(hijax.HiPrim):\n  jvp = hijax.jvp_from_lin\n  lin, linearized = hijax.linearize_from_jvp\n# after (pick one)\nclass P(hijax.HiPrim):\n  def jvp(self, primals, tangents): ...\n  lin, linearized = hijax.linearize_from_jvp","handlingStrategy":"validation","validationCode":"assert not (type(prim).jvp is hijax.jvp_from_lin and\n            type(prim).lin is hijax._lin_from_jvp), 'circular rule derivation'","typeGuard":"def no_circular_rules(p) -> bool:\n    return not (type(p).jvp is hijax.jvp_from_lin and\n                type(p).lin is hijax._lin_from_jvp)","tryCatchPattern":"try:\n    jax.jvp(f, (x,), (t,))\nexcept TypeError as e:\n    if 'defined in terms of the other' in str(e):\n        raise RuntimeError('pick one direction: jvp_from_lin OR linearize_from_jvp') from e\n    raise","preventionTips":["Never combine jvp_from_lin with linearize_from_jvp","Keep a lint/review check on class bodies of HiPrim subclasses"],"tags":["jax","autodiff","circular-definition","misconfiguration"],"backgroundTag":"circular-rule-derivation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}