jax-ml/jax · error · NotImplementedError

callbacks do not support specifying sharding inside spmd com

Error message

callbacks do not support specifying sharding inside spmd computations

What it means

In fully manual SPMD mode, output placement of the callback is already determined by the manual per-device semantics, so an explicitly provided sharding argument is contradictory and unsupported. _callback_op_sharding raises NotImplementedError whenever sharding is not None inside an SPMDAxisContext with all-manual axes.

Source

Thrown at jax/_src/callback.py:163

          mesh_shape=(),
          dim_shardings=(SdyDim(axes=(), is_open=False),) * ndim,
          logical_device_ids=())
      for ndim in ndims))


def _callback_op_sharding(
    axis_context, sharding: Sharding | None, avals_out
):
  if isinstance(axis_context, sharding_impls.SPMDAxisContext):
    # If we have fully manual sharding during lowering, that means the JAX
    # program has per-device semantics, so we run the callback on each device.
    if axis_context.manual_axes != frozenset(axis_context.mesh.axis_names):
      raise NotImplementedError(
          "callbacks are only supported in spmd computations when all mesh"
          " axes are partitioned manually (no partial automatic sharding)."
      )
    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"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Drop the sharding argument when running under manual SPMD (inside shard_map)
Move the callback outside the manual region where specifying sharding is allowed (ShardingContext path supports SingleDeviceSharding)
For multi-device output needs, split outputs and place them after the callback with jax.device_put

Example fix

# before
@shard_map(mesh=mesh, in_specs=P('d'), out_specs=P('d'))
def f(x):
    return pure_callback(fn, x, result_dtype=x.dtype,
                         sharding=NamedSharding(mesh, P('d')))  # raises

# after
@shard_map(mesh=mesh, in_specs=P('d'), out_specs=P('d'))
def f(x):
    return pure_callback(fn, x, result_dtype=x.dtype)  # no sharding arg

Defensive patterns

Strategy: validation

Validate before calling

def spmd_callback(fn, x, sharding):
    if in_manual_spmd():  # e.g. flag set inside shard_map body
        assert sharding is None, 'do not pass sharding under manual SPMD'
    return pure_callback(fn, x, result_dtype=x.dtype, sharding=sharding)

Try / catch

try:
    jax.jit(f)(x)
except NotImplementedError as e:
    if 'do not support specifying sharding inside spmd' in str(e):
        retry with sharding=None
    raise

Prevention

Do not propagate sharding kwargs into callbacks called under shard_map
Make sharding an explicit parameter defaulting to None and omit under manual SPMD
Codify the rule: manual mode owns placement, so never override it

When it happens

Trigger: Calling jax.pure_callback(..., sharding=SomeSharding) or jax.io_callback(..., sharding=...) inside a computation lowered under a Mesh with all axes manual (e.g. inside shard_map or with fully manual partitioning).

Common situations: Passing sharding=NamedSharding(...) to pure_callback while using shard_map for manual control; upgrading code that specified sharding on single-device paths and reusing it in SPMD pipelines.

Related errors

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