jax-ml/jax · error · ValueError

unexpected JAX type (e.g. shape/dtype) for argument to VJP f

Error message

unexpected JAX type (e.g. shape/dtype) for argument to VJP function: got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} because the corresponding output of the differentiated function had JAX type {aval.str_short()}

What it means

In jax.vjp, each cotangent passed to the backward function must have the same shape and dtype (JAX aval) as the corresponding primal output's cotangent aval. JAX raises this when e.g. the cotangent is float64 while the output was float32, or has a different shape.

Source

Thrown at jax/_src/api.py:1977

But the tree structures differ:
"""
  msg += '\n'.join(f"  * out{keystr(path)} was a {thing1} in the original "
                   f"output, but a {thing2} here, so {explanation}."
                   for path, thing1, thing2, explanation
                   in equality_errors_pytreedef(out_tree, ct_tree))
  raise ValueError(msg)


def _vjp_check_ct_avals(cts, primal_avals):
  # TODO(mattjj): improve this error  by flattening with keys in the first place
  for ct, aval in zip(cts, primal_avals):
    if isinstance(ct, ad.Zero): continue
    ct_aval = typeof(ct)
    ct_aval_expected = aval.to_ct_aval()
    if (not core.typecompat(ct_aval, ct_aval_expected) and
        not _temporary_dtype_exception(ct_aval, ct_aval_expected)):
      raise ValueError(
          "unexpected JAX type (e.g. shape/dtype) for argument to VJP function: "
          f"got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} "
          "because the corresponding output of the differentiated function had JAX type "
          f"{aval.str_short()}")


@register_dataclass
@dataclasses.dataclass(frozen=True, slots=True)
class NotNeeded:
  pass

@register_dataclass
@dataclasses.dataclass(frozen=True, slots=True)
class NotSaveable:
  pass

@dataclasses.dataclass(frozen=True, slots=True)
class GradValue:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build cotangents with jax.tree_util.tree_map(jnp.zeros_like, out) or jnp.ones_like
2. Match dtypes explicitly: jnp.ones(shape, dtype=out.dtype)
3. Check jax_enable_x64 consistency between how the function was traced and how cotangents are constructed

Example fix

// before
cts = (np.ones(3.0),)   # float64 vs float32 output
// after
cts = (jnp.ones_like(out_array),)

import jax.tree_util as jtu
cts = jtu.tree_map(lambda o: jnp.zeros(o.shape, o.dtype), out)

try:
    in_cts = vjp_fn(cts)
except ValueError as e:
    if 'unexpected JAX type' in str(e):
        cts = jtu.tree_map(lambda o: jnp.zeros_like(o), out)
        in_cts = vjp_fn(cts)
    else:
        raise

When it happens

Trigger: Passing a numpy float64 array as cotangent for a float32 output; passing a cotangent with wrong shape (e.g. scalar for a (3,) output); mixing jnp and np scalars with different default dtypes.

Common situations: x64 mode enabled (jax.config.update('jax_enable_x64', True)) so np.ones gives float64 while the traced output is float32; constructing cotangents from numpy instead of jnp.zeros_like(out); shape errors from broadcasting assumptions.

Related errors

• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• cotangent type does not match function output, expected {out
• The number of sources must match the packing factor ({packin
• primal and tangent arguments to jax.jvp do not match; dtypes

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