jax-ml/jax · error · RuntimeError

Incorrect output shape for return value #{i}: Expected: {out

Error message

Incorrect output shape for return value #{i}: Expected: {out_aval.shape}, Actual: {out_val.shape}

What it means

Callback output #i has a different shape than the corresponding declared abstract value. JAX checks each returned array's shape against result_avals when the callback executes.

Source

Thrown at jax/_src/callback.py:825

    if result_avals:
      raise ValueError("Partitioned callback not supported with return values.")
  backend: xc.Client = cast(xc.Client, ctx.module_context.get_backend())
  result_shapes = [_aval_to_xla_shape(aval) for aval in result_avals]
  operand_shapes = [_aval_to_xla_shape(aval) for aval in operand_avals]

  # First we apply checks to ensure output shapes and dtypes match the expected
  # ones.
  def _wrapped_callback(*args):
    out_vals = callback(*args)
    if len(out_vals) != len(result_avals):
      raise RuntimeError(
          "Mismatched number of outputs from callback. "
          "Expected: {}, Actual: {}".format(len(result_avals), len(out_vals)))
    # Handle Python literals, and custom arrays, e.g., tf.Tensor.
    out_vals = tuple(dtypes.canonicalize_value(np.asarray(a)) for a in out_vals)
    for i, (out_val, out_aval) in enumerate(zip(out_vals, result_avals)):
      if out_val.shape != out_aval.shape:
        raise RuntimeError(
            f"Incorrect output shape for return value #{i}: "
            f"Expected: {out_aval.shape}, Actual: {out_val.shape}")
      if out_val.dtype != out_aval.dtype:
        raise RuntimeError(
            f"Incorrect output dtype for return value #{i}: "
            f"Expected: {out_aval.dtype}, Actual: {out_val.dtype}")

    if platform == "tpu":
      # On TPU we cannot receive empty arrays. So, we return from the wrapped
      # callback only the non-empty results, and we will create empty constants
      # in the receiving computation.
      # TODO(b/238239458): fix TPU Recv to work with empty arrays.
      non_empty_out_vals = tuple(
          out_val
          for out_val, result_aval in zip(out_vals, result_avals)
          if not is_empty_shape(result_aval.shape))
      return non_empty_out_vals
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Update result_avals/result_shape_dtypes to match the actual output shapes
2. Make the callback reshape/pad its outputs to the declared shape before returning
3. Avoid dynamic output shapes; callbacks must produce the exact declared shape

Example fix

# before
f = lambda x: x.reshape(1, -1)  # declared shape (n,)
pure_callback(f, jax.ShapeDtypeStruct(x.shape, x.dtype), x)
# after
f = lambda x: x
pure_callback(f, jax.ShapeDtypeStruct(x.shape, x.dtype), x)

def make_callback(fn, sample_in):
    out = fn(sample_in)
    avals = tuple(jax.ShapeDtypeStruct(np.shape(o), dtypes.canonicalize_dtype(np.dtype(o))) for o in out)
    return avals
# pass these avals instead of hand-written shapes

try:
    y = jax.pure_callback(fn, avals, x)
except RuntimeError as e:
    if 'Incorrect output shape' in str(e):
        ...log and recompute avals...
    raise

When it happens

Trigger: A callback returning arrays whose shape differs from the shape in result_avals passed to pure_callback/io_callback (e.g. declared (3,) but returned (1,3)).

Common situations: Hardcoded result_shape_dtypes going stale after input shapes change; callbacks producing dynamically sized outputs; rank changes after refactoring.

Related errors

• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output dtype for return value #{i}: Expected: {out
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Buffer callbacks do not support transpose. Please use `jax.c

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