{"record":{"id":"de474fbc32b4826d","repo":"jax-ml/jax","slug":"for-linearize-support-subclass-type-self-must","errorCode":null,"errorMessage":"for linearize support, subclass {type(self)} must implement `lin` and `linearized`, or derive them from its `jvp` rule by setting `lin, linearized = linearize_from_jvp`","messagePattern":"for linearize support, subclass (.+?) must implement `lin` and `linearized`, or derive them from its `jvp` rule by setting `lin, linearized = linearize_from_jvp`","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":194,"sourceCode":"      args_grad, logs = self.vjp_bwd_retval(res, outgrad), None\n    maybe_accum = lambda acc, v: isinstance(acc, ad.GradAccum) and acc.accum(v)\n    tree_map(maybe_accum, arg_accums, args_grad)\n    return logs\n\n  def vjp_bwd_retval(self, res, outgrad, /):\n    # Classic API: returns values instead of using accumulators\n    raise NotImplementedError(\n        f\"for grad support, subclass {type(self)} must implement `vjp_bwd` or \"\n        \"`vjp_bwd_retval`, or derive its reverse-mode rules by setting \"\n        \"`vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= vjp_from_lin`)\")\n\n  # optional forward-mode AD interfaces\n  def jvp(self, primals, tangents):\n    raise NotImplementedError(f\"for jvp support, subclass {type(self)} must \"\n                              \"implement `jvp`\")\n\n  def lin(self, nzs_in, *primals):\n    raise NotImplementedError(\n        f\"for linearize support, subclass {type(self)} must implement `lin` \"\n        \"and `linearized`, or derive them from its `jvp` rule by setting \"\n        \"`lin, linearized = linearize_from_jvp`\")\n\n  def linearized(self, residuals, *tangents):\n    raise NotImplementedError(\n        f\"for linearize support, subclass {type(self)} must implement `lin` \"\n        \"and `linearized`, or derive them from its `jvp` rule by setting \"\n        \"`lin, linearized = linearize_from_jvp`\")\n\n  # optional transpose rule, for primitives that are linear in some inputs\n  def transpose(self, out_ct, *maybe_accums):\n    raise NotImplementedError(f\"for transpose support, subclass {type(self)} \"\n                              \"must implement `transpose`\")\n\n  # vmap interface\n  def batch(self, axis_data, args, dims):\n    out_dim = self.batch_dim_rule(axis_data, dims)","sourceCodeStart":176,"sourceCodeEnd":212,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L176-L212","documentation":"HiPrim's default `lin` rule (linearization/primal part of jax.linearize) is a stub. Subclasses must implement `lin` and `linearized`, or derive both from an existing jvp rule with `lin, linearized = linearize_from_jvp`.","triggerScenarios":"Calling jax.linearize (or code that stages lin, e.g. jvp_from_lin / _vjp_fwd_from_lin derivations) on a HiPrim subclass that lacks linearize rules.","commonSituations":"Setting `vjp_fwd, vjp_bwd_retval = vjp_from_lin` without having lin/linearized defined, or calling jax.linearize on a custom primitive.","solutions":["Set `lin, linearized = linearize_from_jvp` in the class body (requires a jvp rule)","Or implement `def lin(self, nzs_in, *primals)` and `def linearized(self, residuals, *tangents)` directly","Do not combine jvp_from_lin with linearize_from_jvp (circular)"],"exampleFix":"class MyPrim(hijax.HiPrim):\n  def jvp(self, primals, tangents): ...\n  # after: derive linearize\n  lin, linearized = linearize_from_jvp","handlingStrategy":"validation","validationCode":"if type(prim).lin is hijax.HiPrim.lin:\n    raise ValueError(f'{type(prim).__name__} lacks linearize rules')","typeGuard":"def has_lin_rule(p) -> bool:\n    return type(p).lin is not hijax.HiPrim.lin","tryCatchPattern":"try:\n    jax.linearize(f)(x)\nexcept NotImplementedError as e:\n    if 'linearize' in str(e):\n        return jax.jvp(f, (x,), (t,))\n    raise","preventionTips":["Prefer deriving lin/linearized from jvp via linearize_from_jvp","Check rule completeness in a primitive test suite"],"tags":["jax","autodiff","linearize","not-implemented"],"backgroundTag":"missing-autodiff-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}