jax-ml/jax · error · NotImplementedError

pure_callback only supports SingleDeviceSharding, but got {t

Error message

pure_callback only supports SingleDeviceSharding, but got {type(sharding)}

What it means

Outside SPMD (plain ShardingContext lowering), pure_callback lowers via an OpSharding that places computation on a single device; therefore only SingleDeviceSharding (or a scalar NamedSharding, which degenerates to one device) is accepted. Any other sharding type — NamedSharding over a real mesh, PositionalSharding, MultiDeviceSharding — hits NotImplementedError with the offending type name.

Source

Thrown at jax/_src/callback.py:180

    if sharding is not None:
      raise NotImplementedError(
          "callbacks do not support specifying sharding inside spmd"
          " computations"
      )
    if config.use_shardy_partitioner.value:
      op_sharding = _get_sdy_array_list_for_callbacks(avals_out)
    else:
      op_sharding = xc.OpSharding()
      op_sharding.type = xc.OpSharding.Type.MANUAL
    return op_sharding

  if isinstance(axis_context, sharding_impls.ShardingContext):
    if sharding is not None:
      if (isinstance(sharding, sharding_impls.NamedSharding) and
          sharding.mesh.is_scalar):  # pyrefly: ignore[missing-attribute]
        pass
      elif not isinstance(sharding, SingleDeviceSharding):
        raise NotImplementedError(
            "pure_callback only supports SingleDeviceSharding, but got"
            f" {type(sharding)}"
        )
      device = next(iter(sharding.device_set))
      device_assignment = axis_context.device_assignment
      if device_assignment is None:
        raise AssertionError(
            "Please file a bug at https://github.com/jax-ml/jax/issues")
      try:
        device_index = device_assignment.index(device)
      except IndexError as e:
        raise ValueError(
            "Sharding provided to pure_callback specifies a device"
            f" {device} that is not in the device assignment"
            f" ({device_assignment})") from e
    else:
      device_index = 0

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass sharding=None (default; JAX infers a single device) when multi-device sharding is not truly needed
2. Use SingleDeviceSharding for an explicit device: sharding=jax.sharding.SingleDeviceSharding(jax.devices()[0])
3. For multi-device outputs, run the callback per-shard via shard_map (manual SPMD) without an explicit sharding argument

Example fix

# before
sh = NamedSharding(mesh, P('data'))
out = pure_callback(fn, x, result_dtype=x.dtype, sharding=sh)  # NotImplementedError

# after
from jax.sharding import SingleDeviceSharding
out = pure_callback(fn, x, result_dtype=x.dtype,
                    sharding=SingleDeviceSharding(jax.devices()[0]))

from jax.sharding import SingleDeviceSharding, NamedSharding

def ok_sharding(s):
    if s is None: return True
    if isinstance(s, SingleDeviceSharding): return True
    if isinstance(s, NamedSharding) and s.mesh.is_scalar: return True
    return False
assert ok_sharding(sh), f'pure_callback sharding must be single-device, got {type(sharding)}'

from jax.sharding import SingleDeviceSharding, NamedSharding
def is_pure_callback_sharding(s) -> bool:
    return s is None or isinstance(s, SingleDeviceSharding) or (
        isinstance(s, NamedSharding) and bool(s.mesh.is_scalar))

try:
    pure_callback(fn, x, result_dtype=d, sharding=sh)
except NotImplementedError as e:
    if 'only supports SingleDeviceSharding' in str(e):
        sh = SingleDeviceSharding(next(iter(sh.device_set)))
        pure_callback(fn, x, result_dtype=d, sharding=sh)
    raise

When it happens

Trigger: Passing sharding=NamedSharding(mesh, P('data','model')), PositionalSharding, GSPMDSharding, etc. to jax.pure_callback or jax.io_callback in a normal (non-SPMD) jax.jit context.

Common situations: Assuming pure_callback can express multi-device output sharding; reusing a pjit out_shards object as the callback sharding; gradually sharding a pipeline and passing the mesh sharding to a host callback.

Related errors

• callbacks are only supported in spmd computations when all m
• callbacks do not support specifying sharding inside spmd com
• 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/b2449c3ee3bab44c. Report an issue: GitHub.