jax-ml/jax · error · TypeError

at {keystr(path)}, got fwd output type {ty.str_short()} whic

Error message

at {keystr(path)}, got fwd output type {ty.str_short()} which doesn't match primal output type {primal_aval.str_short()}

What it means

During custom_vjp evaluation, JAX checks that each output of the user-supplied forward function has the same type/aval as the corresponding primal output of the @jax.custom_vjp-decorated function. If the fwd rule returns a value whose dtype/shape/sharding differs from the primal's, this TypeError is raised at the mismatching pytree path.

Source

Thrown at jax/_src/hijax.py:890

def _vjp_primal_fwd_tree_mismatch_err(self, tree):
  return (f"Custom VJP fwd rule {self.fwd.__name__} for function {self.traced.fun_name} "
          "must produce a pair (list or tuple of length two) where the first "
          "element represents the primal output "
          "(equal to the output of the custom_vjp-decorated function "
          f"{self.traced.fun_name}) and the "
          "second element represents residuals (i.e. values stored from the "
          "forward pass for use on the backward pass), but "
          f"instead the fwd rule output's first element had container/pytree "
          "structure:\n"
          f"""    {str(tree ).replace("'", "")}\n"""
          f"while the custom_vjp-decorated function {self.traced.fun_name} had output "
          "container/pytree structure:\n"
          f"""    {str(self.out_tree).replace("'", "")}.""")

def _vjp_fwd_aval_mismatch_err(path, primal_aval, fwd_val):
  if not core.typematch(ty := typeof(fwd_val), primal_aval):
    raise TypeError(f"at {keystr(path)}, got fwd output type {ty.str_short()} "
                    f"which doesn't match primal output type {primal_aval.str_short()}")

def _vjp_bwd_aval_mismatch_err(primal_sourceinfo, path, primal_aval, ct):
  if config.disable_bwd_checks.value:
    return
  if isinstance(ct, ad_util.Zero):
    return
  if isinstance(primal_aval, AbstractRef):
    primal_aval = primal_aval.inner_aval
  expected = primal_aval.to_ct_aval()
  ct_aval = ct.aval if isinstance(ct, ad_util.SymbolicZero) else typeof(ct)
  if (not core.typematch(expected, ct_aval) and
      not _temporary_dtype_exception(expected, ct_aval) and
      getattr(expected, 'dtype', None) is not dtypes.float0):
    result = f"at output{keystr(path)} " if path else ""
    raise ValueError(
        f"{result}the bwd rule attached to {primal_sourceinfo} produced an"
        f" output of type {ct_aval.str_short()} which doesn't match expected"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the fwd rule return the primal outputs with exactly the same avals (shape+dtype) as the decorated function
2. Check for implicit dtype casts in the fwd function (e.g. Python floats mixing with float32, x64 mismatches) and cast explicitly with jax.lax.convert_element_type to the primal dtype
3. Ensure the pytree structure returned by fwd matches the decorated function's output structure exactly

Example fix

// before
@jax.custom_vjp
def f(x):
  return x.astype(jnp.float32) * 2
def f_fwd(x):
  return x * 2, None          # returns float64 if x is float64
// after
def f_fwd(x):
  primal = x.astype(jnp.float32) * 2   # match primal aval exactly
  return primal, None

# verify fwd outputs match primal avals before use
import jax
flat_primal = jax.eval_shape(f, x)
flat_fwd, _ = jax.eval_shape(lambda x: (f_fwd(x)), x)  # first of pair
assert jax.tree_util.tree_structure(flat_fwd) == jax.tree_util.tree_structure(flat_primal)

try:
    y = f(x)
except TypeError as e:
    if 'fwd output type' in str(e):
        # align fwd return avals with primal outputs and retry
        ...

When it happens

Trigger: Passing a fwd function to custom_vjp.defvjp whose return values (e.g. residuals or primals) differ in shape or dtype from what the decorated function returns, e.g. returning float32 residuals when the primal outputs float64, or returning a differently nested pytree.

Common situations: Writing a custom_vjp fwd rule that casts or reshapes outputs (e.g. l.to(torch.float32) style casts, or int truncation) inconsistently with the primal function; changing the primal signature without updating the fwd rule; enabling x64 mode after writing rules with hardcoded dtypes.

Related errors

• {result}the bwd rule attached to {primal_sourceinfo} produce
• Pure callbacks do not support transpose. Please use `jax.cus
• No VJP defined for custom_vjp function {self.f.__name__} usi
• custom_vjp symbolic_zeros support with shard_map is not impl
• primal and tangent arguments to jax.jvp must be tuples or li

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