jax-ml/jax · error · TypeError

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

Error message

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

What it means

vjp_from_jvp is a NamedTuple pair (vjp_fwd, vjp_bwd_retval); it is not itself callable. Invoking it directly raises this TypeError pointing to the required class-body unpacking.

Source

Thrown at jax/_src/hijax.py:671

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Unpack in the class body: `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`

Example fix

# before
class P(hijax.HiPrim):
  vjp_fwd = vjp_from_jvp
# after
class P(hijax.HiPrim):
  vjp_fwd, vjp_bwd_retval = vjp_from_jvp

from jax._src import hijax
for attr in ('vjp_fwd', 'vjp_bwd_retval'):
    assert not isinstance(type(p).__dict__.get(attr), hijax._VJPFromJVP), attr + ' not unpacked'

def vjp_rules_valid(p) -> bool:
    from jax._src import hijax
    d = type(p).__dict__
    return not isinstance(d.get('vjp_fwd'), hijax._VJPFromJVP) and \
           not isinstance(d.get('vjp_bwd_retval'), hijax._VJPFromJVP) and \
           not isinstance(d.get('vjp_fwd'), 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_jvp') from e
    raise

When it happens

Trigger: Writing `vjp_fwd = vjp_from_jvp` or calling `vjp_from_jvp(res, ct)` instead of `vjp_fwd, vjp_bwd_retval = vjp_from_jvp`.

Common situations: Migrating from classic custom_vjp where a single bwd function is registered.

Related errors

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