jax-ml/jax · error · TypeError

`vjp_from_lin` is a pair of rules, not a single rule; unpack

Error message

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

What it means

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.

Source

Thrown at jax/_src/hijax.py:679

  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`")

linearize_from_jvp = _LinearizeFromJVP(_lin_from_jvp, _linearized_from_jvp)
vjp_from_jvp = _VJPFromJVP(_vjp_fwd_from_jvp, _transpose_jvp)
vjp_from_lin = _VJPFromLin(_vjp_fwd_from_lin, _transpose_linearized)


class CustomVJPTraced(HiPrim):
  """Applications take ``(consts, fwd_consts, *args)``.

  The two leading arguments are synthetic, and both get zero cotangents:
  ``consts`` is the primal function's closure environment promoted to an
  argument (see ``Traced.with_consts_as_arg``), and ``fwd_consts`` is extra
  inputs consumed only by the fwd rule and ignored by the primal (ordinarily
  ``()``; the ``remat`` rule uses it to pass replay residuals to the helper
  primitive it builds). Any value a rule needs beyond the primal arguments
  must arrive through these slots as an explicit input, never by closure: the

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Unpack in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_lin` (requires lin/linearized to be defined, not derived from jvp)

Example fix

# before
class P(hijax.HiPrim):
  vjp_bwd_retval = vjp_from_lin
# after
class P(hijax.HiPrim):
  lin, _ = linearize_from_jvp
  vjp_fwd, vjp_bwd_retval = vjp_from_lin

from jax._src import hijax
assert not isinstance(type(p).__dict__.get('vjp_fwd'), hijax._VJPFromLin)

def vjp_from_lin_valid(p) -> bool:
    from jax._src import hijax
    return not isinstance(type(p).__dict__.get('vjp_fwd'), hijax._VJPFromLin) and \
           not isinstance(type(p).__dict__.get('vjp_bwd_retval'), hijax._VJPFromLin)

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'pair of rules' in str(e):
        raise RuntimeError('unpack: vjp_fwd, vjp_bwd_retval = vjp_from_lin') from e
    raise

When it happens

Trigger: Writing `vjp_fwd = vjp_from_lin` or calling `vjp_from_lin(...)` directly rather than unpacking `vjp_fwd, vjp_bwd_retval = vjp_from_lin`.

Common situations: 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.

Related errors

• `linearize_from_jvp` is a pair of rules, not a single rule;
• `vjp_from_jvp` 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/9688aa4f186eca31. Report an issue: GitHub.