jax-ml/jax · error · TypeError

cotangent type does not match function output, expected {out

Error message

cotangent type does not match function output, expected {out_avals} but got {out_cts}

What it means

Even when the pytree structure matches, each cotangent leaf passed to a jax.linear_transpose pullback must typecheck against the corresponding output aval (shape and dtype). This fires on shape or dtype mismatches per leaf.

Source

Thrown at jax/_src/api.py:2124

                                                      instantiate=True)
  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)
  out_avals, _ = unzip2(out_pvals)
  out_dtypes = [a.dtype for a in out_avals if not a.is_high]
  if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)
          or all(dtypes.issubdtype(d, np.integer)
                 for d in in_dtypes + out_dtypes)):
    raise TypeError("linear_transpose only supports [float or complex] -> "
                    "[float or complex], and integer -> integer functions, "
                    f"but got {in_dtypes} -> {out_dtypes}.")

  @api_boundary
  def transposed_fun(const, out_cotangent):
    out_cts, out_tree2 = tree_flatten(out_cotangent)
    if out_tree() != out_tree2:
      raise TypeError("cotangent tree does not match function output, "
                      f"expected {out_tree()} but got {out_tree2}")
    if not all(map(core.typecheck, out_avals, out_cts)):
      raise TypeError("cotangent type does not match function output, "
                      f"expected {out_avals} but got {out_cts}")
    dummies = [ad.UndefinedPrimal(a.to_ct_aval()) for a in in_avals]
    in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts)
    in_cts = map(ad.instantiate_zeros, in_cts)
    return tree_unflatten(in_tree, in_cts)

  # Ensure that transposed_fun is a PyTree
  return Partial(transposed_fun, const)


@overload
def make_jaxpr(
    fun: Callable,
    static_argnums: int | Sequence[int] = (),
    axis_env: Sequence[tuple[AxisName, int]] | None = None,
    return_shape: Literal[False] = ...,
) -> Callable[..., core.Jaxpr]:
  ...

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct cotangents with jnp.ones_like / jnp.zeros_like on the actual outputs
2. Specify dtype explicitly: jnp.zeros(shape, dtype=out.dtype)
3. Ensure jax_enable_x64 setting matches between tracing and cotangent construction

Example fix

# before
ct = np.zeros(x.shape)          # possibly float64
# after
ct = jnp.zeros_like(f(x))

import jax.tree_util as jtu
out = f(*primals)
ct = jtu.tree_map(lambda o: jnp.zeros(o.shape, o.dtype), out)  # dtype/shape-exact

try:
    in_ct = f_t(ct)
except TypeError as e:
    if 'cotangent type' in str(e):
        ct = jtu.tree_map(jnp.zeros_like, f(*primals))
        in_ct = f_t(ct)
    else:
        raise

When it happens

Trigger: Passing a float64 cotangent for a float32 output; passing shape (2,3) where output was (3,2); passing a Python scalar where a tracer-compatible array is required.

Common situations: x64 enabled causing numpy-built cotangents to be float64; reshaping/broadcast mistakes; passing zeros with an inferred dtype that differs from the output dtype.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• linear_transpose only supports [float or complex] -> [float
• The number of sources must match the packing factor ({packin
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass

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