jax-ml/jax · error · NotImplementedError

for linearize support, subclass {type(self)} must implement

Error message

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`

What it means

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`.

Source

Thrown at jax/_src/hijax.py:194

      args_grad, logs = self.vjp_bwd_retval(res, outgrad), None
    maybe_accum = lambda acc, v: isinstance(acc, ad.GradAccum) and acc.accum(v)
    tree_map(maybe_accum, arg_accums, args_grad)
    return logs

  def vjp_bwd_retval(self, res, outgrad, /):
    # Classic API: returns values instead of using accumulators
    raise NotImplementedError(
        f"for grad support, subclass {type(self)} must implement `vjp_bwd` or "
        "`vjp_bwd_retval`, or derive its reverse-mode rules by setting "
        "`vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= vjp_from_lin`)")

  # optional forward-mode AD interfaces
  def jvp(self, primals, tangents):
    raise NotImplementedError(f"for jvp support, subclass {type(self)} must "
                              "implement `jvp`")

  def lin(self, nzs_in, *primals):
    raise NotImplementedError(
        f"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`")

  def linearized(self, residuals, *tangents):
    raise NotImplementedError(
        f"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`")

  # optional transpose rule, for primitives that are linear in some inputs
  def transpose(self, out_ct, *maybe_accums):
    raise NotImplementedError(f"for transpose support, subclass {type(self)} "
                              "must implement `transpose`")

  # vmap interface
  def batch(self, axis_data, args, dims):
    out_dim = self.batch_dim_rule(axis_data, dims)

View on GitHub (pinned to 1e1c6a8fc0)

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)

Example fix

class MyPrim(hijax.HiPrim):
  def jvp(self, primals, tangents): ...
  # after: derive linearize
  lin, linearized = linearize_from_jvp

Defensive patterns

Strategy: validation

Validate before calling

if type(prim).lin is hijax.HiPrim.lin:
    raise ValueError(f'{type(prim).__name__} lacks linearize rules')

Type guard

def has_lin_rule(p) -> bool:
    return type(p).lin is not hijax.HiPrim.lin

Try / catch

try:
    jax.linearize(f)(x)
except NotImplementedError as e:
    if 'linearize' in str(e):
        return jax.jvp(f, (x,), (t,))
    raise

Prevention

Prefer deriving lin/linearized from jvp via linearize_from_jvp
Check rule completeness in a primitive test suite

When it happens

Trigger: 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.

Common situations: Setting `vjp_fwd, vjp_bwd_retval = vjp_from_lin` without having lin/linearized defined, or calling jax.linearize on a custom primitive.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/de474fbc32b4826d. Report an issue: GitHub.