jax-ml/jax · error · ValueError

Partitioned callback not supported with return values.

Error message

Partitioned callback not supported with return values.

What it means

A partitioned callback was requested but the callback declares result values (result_avals). Partitioned callbacks must be side-effecting procedures with no return values.

Source

Thrown at jax/_src/callback.py:808

      False, then `callback` is called on all shards.
    sharding: The sharding of the callback.

  Returns:
    A tuple of MLIR result values, a new token (if any), and the host callback
    object.
  """
  if len(ctx.module_context.platforms) > 1:
    raise NotImplementedError("multi-platform lowering for python_callback")
  platform = ctx.module_context.platforms[0]
  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}: "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove return values from the callback and return nothing (make it a side-effecting callback)
2. If you need results, use a non-partitioned pure_callback/io_callback

Example fix

# before
io_callback(lambda x: x * 2, make_jaxpr(...).out_avals, x, _partitioned=True)
# after
io_callback(lambda x: None, (), x, _partitioned=True)  # side effects only

def make_partitioned_callback(fn, result_avals):
    if result_avals:
        raise ValueError('use non-partitioned callback for value-returning fns')
    return partial(io_callback, fn, result_avals, _partitioned=True)

When it happens

Trigger: Calling io_callback/pure_callback with _partitioned=True while passing a non-empty result_avals tuple, i.e. the callback returns values.

Common situations: Adapting existing pure_callback code (which returns values) to run partitioned; misunderstanding that partitioned callbacks are emit-only.

Related errors

• Partitioned callback not implemented on {platform} backend.
• numpy masked arrays are not supported as direct inputs to JA
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• to_dlpack can only pack a dlpack tensor from an array on a s

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