jax-ml/jax · error · NotImplementedError

Partitioned callback not implemented on {platform} backend.

Error message

Partitioned callback not implemented on {platform} backend.

What it means

Raised during lowering of a partitioned JAX callback (io_callback/pure_callback with _partitioned=True) on a backend that doesn't support partitioning. Only cpu, cuda, rocm, and oneapi support partitioned callbacks; TPU does not.

Source

Thrown at jax/_src/callback.py:805

    has_side_effect: Whether the callback has side effects.
    returns_token: Whether the callback should return a token.
    partitioned: If True, then `callback` is called on local shards only. If
      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)):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch to a non-partitioned callback: drop _partitioned=True and use a plain jax.experimental.io_callback / pure_callback
2. Run on a supported backend (cpu/cuda/rocm/oneapi) if partitioning is required
3. Upgrade JAX — check release notes for added TPU partitioned-callback support

Example fix

# before
io_callback(fn, result_avals, *args, _partitioned=True)  # on TPU
# after
io_callback(fn, result_avals, *args)  # unpartitioned, works on TPU

from jax._src.callback import _PARTITIONED_PLATFORMS  # illustrative
supported = {'cpu','cuda','rocm','oneapi'}
platform = jax.default_backend()
if platform not in supported:
    callback_kwargs.pop('_partitioned', None)  # degrade gracefully

When it happens

Trigger: Calling io_callback(..., _partitioned=True) (as done inside jax.debug, pallas, or TPU-scheduled code) while compiling/executing on the TPU backend (or any platform outside the supported set).

Common situations: Running code that was written for GPU partitioned callbacks on a TPU device or cloud TPU VM; using libraries (e.g. extendedblas, pallas TPU) that internally request partitioned callbacks.

Related errors

• Partitioned callback not supported with return values.
• Custom Partitioning rules must return Sharding.
• The SVD algorithm parameter is not implemented on TPU.
• Mosaic kernels cannot be automatically partitioned. Please w
• The sharding must divide the mask blocks evenly between devi

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