jax-ml/jax · error · RuntimeError

tree mismatch during linearization of {prim=}. Expected: {pr

Error message

tree mismatch during linearization of {prim=}. Expected: {prim.out_tree} got: {treedef_actual}

What it means

The user-supplied `linearized` rule of a HiPrim returned tangents whose pytree structure doesn't match the primitive's declared out_tree. flatten_user_linearized flattens the rule output and compares treedefs, raising RuntimeError on mismatch.

Source

Thrown at jax/_src/hijax.py:456

  nz_tangents = tree_leaves(tangents)
  out_nz = call_hi_primitive_linearized_p.bind(
      *residuals_flat, *nz_tangents, residuals_tree=residuals_tree, _prim=prim,
      nz_in_flat=tuple(nz_in_flat), nz_out_flat=tuple(nz_out_flat),
      has_sres=sres is not None)
  out_nz_iter = iter(out_nz)
  out = [next(out_nz_iter) if nz else ad_util.Zero(a.to_tangent_aval())
         for a, nz in zip(prim.out_avals_flat, nz_out_flat)]
  assert next(out_nz_iter, sentinel := object()) is sentinel
  return out

def flatten_user_linearized(prim, residuals, sres, *tangents_flat):
  tangents = tree_unflatten(prim.in_tree, tangents_flat)
  tangents_out = (prim.linearized(residuals, *tangents) if sres is None else
                  prim.linearized(residuals, sres, *tangents))
  flat_vals, treedef_actual = tracing_registry.flatten(
      tangents_out, lambda x: isinstance(x, ad_util.Zero))
  if treedef_actual != prim.out_tree:
    raise RuntimeError(
        f"tree mismatch during linearization of {prim=}."
        f" Expected: {prim.out_tree} got: {treedef_actual}"
    )
  return flat_vals

call_hi_primitive_linearized_p = core.Primitive("call_hi_primitive_linearized")
call_hi_primitive_linearized_p.multiple_results = True
call_hi_primitive_linearized_p.is_high = lambda *args, _prim, **_: True
@call_hi_primitive_linearized_p.def_abstract_eval
def _call_hi_primitive_linearized_abstract_eval(
    *_args, _prim, residuals_tree, nz_in_flat, nz_out_flat, has_sres):
  return [t.to_tangent_aval() for t, nz in zip(_prim.out_avals_flat, nz_out_flat) if nz]

def _call_hi_primitive_linearized_transpose(
    cts_flat_, *args, _prim, residuals_tree, nz_in_flat, nz_out_flat, has_sres):
  residuals_flat, accums_flat = split_list(args, [residuals_tree.num_leaves])
  residuals = tree_unflatten(residuals_tree, residuals_flat)
  accums_flat_ = iter(accums_flat)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make linearized return exactly the same pytree structure as the primal output (mirror out_tree)
2. Use jax.tree_util.tree_unflatten(prim.out_tree, ...) or build the output the same way the primal does
3. Add a unit test comparing tree_structure(prim.linearized(...)) with tree_structure(primal output)

Example fix

# before
def linearized(self, residuals, *ts):
    return [ts[0]]  # wrong structure
# after
def linearized(self, residuals, *ts):
    return {'y': ts[0]}  # matches out_tree dict

from jax.tree_util import tree_structure
assert tree_structure(dummy_out) == tree_structure(prim.out_tree.unflatten([0]*n))

def linearized_tree_ok(prim, residuals, tangents) -> bool:
    from jax.tree_util import tree_structure
    return tree_structure(prim.linearized(residuals, *tangents)) == prim.out_tree

try:
    f_lin(t)
except RuntimeError as e:
    if 'tree mismatch during linearization' in str(e):
        raise RuntimeError('fix linearized return structure to match primal out_tree') from e
    raise

When it happens

Trigger: A custom linearized(residuals, *tangents) implementation returning e.g. a tuple instead of the dict, wrong arity, or differently-nested containers than the primal output tree.

Common situations: Hand-written linearize rules that forget an output component or wrap results in an extra list/tuple.

Related errors

• primal and tangent arguments to jax.jvp must have the same t
• expected {in_avals.tree}, got {tangents_ft.tree}
• structure of the differentiated function {jaxpr.debug_info.f
• Custom JVP rule {jvp_name} for function {primal_name} must p
• Custom JVP rule {jvp_name} for function {primal_name} must p

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