jax-ml/jax · error · TypeError

subclass {type(self)} can't set both `jvp = jvp_from_lin` an

Error message

subclass {type(self)} can't set both `jvp = jvp_from_lin` and `lin, linearized = linearize_from_jvp`, since each would be defined in terms of the other

What it means

jvp_from_lin detects that the class also set `lin, linearized = linearize_from_jvp`, i.e. both forward rules would be mutually defined in terms of each other, causing infinite recursion; it raises TypeError immediately.

Source

Thrown at jax/_src/hijax.py:603

    return out_primals_flat, out_tangents_flat

  dbg = debug_info('linearize_from_jvp', self.jvp, (primals, primals), {})
  out_primals_flat, nzs_out_flat, consts, _, linearized = ad.linearize_from_jvp(
      lu.wrap_init(jvp_flat, debug_info=dbg), True, nzs_in_flat,
      False, False, primals_flat, {})
  out_primals = tree_unflatten(self.out_tree, out_primals_flat)
  nzs_out = tree_unflatten(self.out_tree, list(nzs_out_flat))
  return out_primals, DerivedLinearization(consts, linearized), nzs_out

def _linearized_from_jvp(self, residuals, *tangents):
  """The `linearized` half of the `linearize_from_jvp` pair."""
  tangents_flat = self.in_tree.flatten_up_to(tangents)
  out_tangents_flat = residuals.apply(residuals.consts, None, *tangents_flat)
  return tree_unflatten(self.out_tree, out_tangents_flat)

def jvp_from_lin(self, primals, tangents):
  if type(self).lin is _lin_from_jvp:
    raise TypeError(
        f"subclass {type(self)} can't set both `jvp = jvp_from_lin` and "
        "`lin, linearized = linearize_from_jvp`, since each would be defined "
        "in terms of the other")
  tangents_flat = self.in_tree.flatten_up_to(tangents)
  nzs_in = tree_unflatten(
      self.in_tree, [not isinstance(t, ad_util.Zero) for t in tangents_flat])
  out_primals, residuals, *rest = self.lin(nzs_in, *primals)
  out_tangents = (self.linearized(residuals, *tangents) if len(rest) < 2 else
                  self.linearized(residuals, rest[1], *tangents))
  return out_primals, out_tangents

def _vjp_fwd_from_jvp(self, nzs_in, *primals):
  """The `vjp_fwd` half of the `vjp_from_jvp` pair."""
  return self(*primals), (primals, nzs_in)

def _transpose_jvp(self, res, out_ct):
  """The `vjp_bwd_retval` half of the `vjp_from_jvp` pair."""
  primals, nzs_in = res

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep exactly one direction: either define jvp and use linearize_from_jvp, or define lin/linearized and use jvp_from_lin
2. Remove `jvp = jvp_from_lin` if you have real lin/linearized rules
3. Remove `lin, linearized = linearize_from_jvp` if you have a real jvp rule

Example fix

# before
class P(hijax.HiPrim):
  jvp = hijax.jvp_from_lin
  lin, linearized = hijax.linearize_from_jvp
# after (pick one)
class P(hijax.HiPrim):
  def jvp(self, primals, tangents): ...
  lin, linearized = hijax.linearize_from_jvp

assert not (type(prim).jvp is hijax.jvp_from_lin and
            type(prim).lin is hijax._lin_from_jvp), 'circular rule derivation'

def no_circular_rules(p) -> bool:
    return not (type(p).jvp is hijax.jvp_from_lin and
                type(p).lin is hijax._lin_from_jvp)

try:
    jax.jvp(f, (x,), (t,))
except TypeError as e:
    if 'defined in terms of the other' in str(e):
        raise RuntimeError('pick one direction: jvp_from_lin OR linearize_from_jvp') from e
    raise

When it happens

Trigger: In a HiPrim class body setting both `jvp = jvp_from_lin` and `lin, linearized = linearize_from_jvp`.

Common situations: Copy-pasting all derivation decorators from the hijax docs without realizing these two are alternative directions (jvp-from-linearize vs linearize-from-jvp).

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which

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