jax-ml/jax · error · TypeError

`linearize_from_jvp` is a pair of rules, not a single rule;

Error message

`linearize_from_jvp` is a pair of rules, not a single rule; unpack it in the class body: `lin, linearized = linearize_from_jvp`

What it means

linearize_from_jvp is a NamedTuple pair of rules (_lin_from_jvp, _linearized_from_jvp), not a callable rule. Calling it directly raises this TypeError telling you to unpack it in the class body.

Source

Thrown at jax/_src/hijax.py:663

def _vjp_fwd_from_lin(self, nzs_in, *primals):
  """The `vjp_fwd` half of the `vjp_from_lin` pair."""
  return self.lin(nzs_in, *primals)

def _transpose_linearized(self, residuals, out_ct):
  """The `vjp_bwd_retval` half of the `vjp_from_lin` pair."""
  def tangent_map(*tangents):
    return self.linearized(residuals, *tangents)
  zero = lambda x: isinstance(x, ad_util.Zero)
  out_ct = tree_map(ad_util.instantiate, out_ct, is_leaf=zero)
  dummies = tree_map(lambda a: ad_util.zeros_like_aval(a.to_tangent_aval()),
                     self.in_avals)
  return api.linear_transpose(tangent_map, *dummies)(out_ct)

class _LinearizeFromJVP(NamedTuple):
  lin: Callable
  linearized: Callable
  def __call__(self, *args, **kwargs):
    raise TypeError(
        "`linearize_from_jvp` is a pair of rules, not a single rule; unpack "
        "it in the class body: `lin, linearized = linearize_from_jvp`")

class _VJPFromJVP(NamedTuple):
  vjp_fwd: Callable
  vjp_bwd_retval: Callable
  def __call__(self, *args, **kwargs):
    raise TypeError(
        "`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`")

class _VJPFromLin(NamedTuple):
  vjp_fwd: Callable
  vjp_bwd_retval: Callable
  def __call__(self, *args, **kwargs):
    raise TypeError(
        "`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`")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Unpack in the class body: `lin, linearized = linearize_from_jvp`
2. If only lin is needed, use `lin, _ = linearize_from_jvp`

Example fix

# before
class P(hijax.HiPrim):
  lin = linearize_from_jvp  # wrong
# after
class P(hijax.HiPrim):
  lin, linearized = linearize_from_jvp

from jax._src import hijax
assert callable(getattr(prim, 'lin', None)) and not isinstance(getattr(type(prim), 'lin', None), hijax._LinearizeFromJVP)

def lin_rule_valid(p) -> bool:
    from jax._src import hijax
    lin = type(p).__dict__.get('lin')
    return lin is not None and not isinstance(lin, hijax._LinearizeFromJVP)

try:
    jax.linearize(f, x)
except TypeError as e:
    if 'pair of rules' in str(e):
        raise RuntimeError('unpack: lin, linearized = linearize_from_jvp') from e
    raise

When it happens

Trigger: Assigning `lin = linearize_from_jvp` or calling `linearize_from_jvp(...)` directly instead of unpacking `lin, linearized = linearize_from_jvp`.

Common situations: Assuming the module-level constant is a single function; using it with defvjp-style single-rule conventions.

Related errors

• `vjp_from_jvp` is a pair of rules, not a single rule; unpack
• `vjp_from_lin` is a pair of rules, not a single rule; unpack
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li
• check_error takes an Error as argument, got type {type(error

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