jax-ml/jax · error · RuntimeError

Mismatched number of outputs from callback. Expected: {}, Ac

Error message

Mismatched number of outputs from callback. Expected: {}, Actual: {}

What it means

The wrapped Python callback returned a different number of outputs than declared in result_avals. JAX validates callback results at runtime against the promised abstract values.

Source

Thrown at jax/_src/callback.py:818

  if platform not in {"cpu", "cuda", "rocm", "tpu", "oneapi"}:
    raise ValueError(
        f"`EmitPythonCallback` not supported on {platform} backend.")
  if partitioned:
    if platform not in {"cpu", "cuda", "rocm", "oneapi"}:
      raise NotImplementedError(
          f"Partitioned callback not implemented on {platform} backend.")
    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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the callback return exactly as many outputs as declared in result_avals (wrap single values in a tuple)
2. Re-check the result_shape_dtypes declaration matches the function's actual return
3. Return () explicitly for zero-result callbacks instead of None

Example fix

# before
f = lambda x: (x + 1)  # declared 2 results
pure_callback(f, (shapes, shapes2), x)
# after
f = lambda x: (x + 1, x + 2)
pure_callback(f, (shapes, shapes2), x)

def checked_pure_callback(fn, result_avals, *args):
    out = fn(*args)  # dry-run on sample inputs in tests
    assert len(jtu.tree_leaves(out)) == len(result_avals), (len(out), len(result_avals))
    return jax.pure_callback(fn, result_avals, *args)

try:
    y = jax.pure_callback(fn, avals, x)
except RuntimeError as e:
    if 'Mismatched number of outputs' in str(e):
        raise ValueError('callback arity drift; re-sync result_avals') from e
    raise

When it happens

Trigger: A pure_callback/io_callback whose wrapped function returns a tuple of length different from len(result_avals); e.g. returning a single array when two were declared, or returning None.

Common situations: Callback returns a scalar instead of a 1-tuple; callback branches and returns different tuple sizes; result_shape/dtype declaration out of sync with the function after refactoring.

Related errors

• Incorrect output shape for return value #{i}: Expected: {out
• 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/1b6e2a7fbdd7a9a8. Report an issue: GitHub.