jax-ml/jax · error · ValueError

linearized function called on tangent values inconsistent wi

Error message

linearized function called on tangent values inconsistent with the original primal values:
Got tangent aval {tangent_aval} for primal aval {primal_aval} but expected {expected_tangent_aval}.{extra_msg}

What it means

The linearized callable checks each tangent leaf's abstract value (shape/dtype, including varying/precision qualifiers) against the tangent aval expected from the original primal. Mismatches raise this, with suggestions (e.g. applying jax.lax.pcast with to='varying') when relevant.

Source

Thrown at jax/_src/api.py:1602

    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"
              " the tangent value passed to the callable `f_jvp` returned by"
              " `jax.linearize`")
        extra_msg = " \nThis might be fixed by:\n" + "\n".join(
            f"  * {d};" for d in pvary_applications)
      raise ValueError(
          "linearized function called on tangent values inconsistent with "
          "the original primal values:\n"
          f"Got tangent aval {tangent_aval} for primal aval {primal_aval} "
          f"but expected {expected_tangent_aval}.{extra_msg}")
  sres_flat = tree_leaves(structured_residuals)
  tangents_out = eval_jaxpr(jaxpr, consts, *tangents_ft, *sres_flat)
  tangents_out_ = iter(tangents_out)
  full_out = [a2tz(aval).instantiate() if known else next(tangents_out_)
              for aval, known in zip(out_avals, out_zeros)]
  assert next(tangents_out_, None) is None
  return out_avals.update(full_out).unflatten()

# TODO(mattjj): see similar function in custom_derivatives.py
def _temporary_dtype_exception(a, a_) -> bool:
  if isinstance(a, core.ShapedArray) and isinstance(a_, core.ShapedArray):
    return a.shape == a_.shape and a_.dtype == float0
  return False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Follow the extra_msg hints: apply jax.lax.pcast(t, to='varying') to tangent values when primals are varying
2. Rebuild tangents from the exact primals used at linearize time (zeros_like/instantiate)
3. Ensure shape/dtype/precision context matches between linearize and the call

Example fix

# before
f_lin = jax.linearize(f)(x)
y = f_lin(jnp.ones_like(x))
# after (varying-precision primals)
y = f_lin(jax.lax.pcast(jnp.ones_like(x), (), to='varying'))

for pa, t in zip(in_avals, tree_leaves(tangents)):
    exp = pa.to_tangent_aval()
    assert core.typecompat(exp, jax.core.typeof(t)), f'{t} incompatible with {exp}'

def tangent_aval_ok(primal_aval, t): return core.typecompat(primal_aval.to_tangent_aval(), jax.core.typeof(t))

When it happens

Trigger: Calling jax.linearize(f)(...) result with tangents of wrong dtype/shape, or with non-'varying' values where the primal was varying-precision and pcast is required on tangents.

Common situations: Mixed-precision (varying) pipelines where tangents must be marked varying via jax.lax.pcast; changing precision context between linearize and the later call; reusing tangents from a differently-shaped input.

Related errors

• expected {in_avals.tree}, got {tangents_ft.tree}
• {self!r} has no `sharding`.
• Custom JVP rule must produce primal and tangent outputs with
• Custom JVP rule must produce primal and tangent outputs with
• for linearize support, subclass {type(self)} must implement

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