jax-ml/jax · error · TypeError

expected {in_avals.tree}, got {tangents_ft.tree}

Error message

expected {in_avals.tree}, got {tangents_ft.tree}

What it means

When calling the function returned by jax.linearize, the tangents passed in must match the pytree structure of the original primal arguments. This TypeError reports the expected vs got tree structures.

Source

Thrown at jax/_src/api.py:1579

  primals_ft = ft.flatten(primals)
  in_nzs_flat = None if in_nzs is None else tuptree_flags(
      in_nzs, primals_ft.tree, 'in_nzs', 'the in_nzs argument to jax.linearize')
  out_primals_ft, out_zeros, jaxpr, consts, structured_residuals, *maybe_aux = \
      ad.linearize(fun, primals_ft, has_aux=has_aux, in_nzs=in_nzs_flat)
  in_avals = primals_ft.map(core.typeof)
  out_avals = out_primals_ft.map(core.typeof)
  lifted_jvp = Partial(
      partial(_lift_linearized, jaxpr, in_avals, out_avals, out_zeros),
      consts, structured_residuals)
  lifted_jvp.out_nzs = tuple(not z for z in out_zeros)  # pyrefly: ignore[missing-attribute]
  return out_primals_ft.unflatten(), lifted_jvp, *maybe_aux


def _lift_linearized(jaxpr, in_avals, out_avals, out_zeros, consts,
                     structured_residuals, *tangents):
  tangents_ft = ft.flatten(tangents)
  if tangents_ft.tree != in_avals.tree:
    raise TypeError(f"expected {in_avals.tree}, got {tangents_ft.tree}")

  tangent_avals = tangents_ft.map(core.typeof)
  for primal_aval, tangent_aval in zip(in_avals, tangent_avals):
    expected_tangent_aval  = primal_aval.to_tangent_aval()
    if not core.typecompat(expected_tangent_aval, tangent_aval):
      extra_msg = ''
      if (isinstance(primal_aval, core.ShapedArray) and
          isinstance(tangent_aval, core.ShapedArray) and
          primal_aval.mat != tangent_aval.mat):
        # TODO(yashkatariya): Tweak error.
        pvary_applications = []
        if left := tangent_aval.mat.varying - primal_aval.mat.varying:
          pvary_applications.append(
              f"applying `jax.lax.pcast(..., {tuple(left)}, to='varying')` to"
              " the primal value passed to `jax.linearize`")
        if left := primal_aval.mat.varying - tangent_aval.mat.varying:
          pvary_applications.append(
              f"applying `jax.lax.pcast(..., {tuple(left)}, to='varying')` to"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass tangents structured exactly like the original primals (reuse the primal pytree via tree.map)
2. Capture the original tree and rebuild tangents with jax.tree.unflatten before calling
3. Write a small wrapper that flattens inputs the same way as at linearize time

Example fix

# before
f_lin = jax.linearize(f)(x, y)
out = f_lin(jnp.ones_like(x))  # missing y
# after
f_lin = jax.linearize(f)(x, y)
out = f_lin(jnp.ones_like(x), jnp.zeros_like(y))

assert jax.tree.structure(tangents) == jax.tree.structure(original_primals), 'linearize tangents must mirror primal tree'

def tangent_tree_ok(primals, t): return jax.tree.structure(t) == jax.tree.structure(primals)

When it happens

Trigger: linearize(f)(x)(t) where t has different structure than x (e.g. x was a tuple (a, b) but a flat tangent or differently nested structure is passed).

Common situations: Saving/reloading a linearized function and calling it with differently-packed inputs; partial-application refactor changing arity; passing cotangents from a vjp (output structure) into a linearized call (input structure).

Related errors

• primal and tangent arguments to jax.jvp must have the same t
• tree mismatch during linearization of {prim=}. Expected: {pr
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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