jax-ml/jax · error · NotImplementedError

callbacks are only supported in spmd computations when all m

Error message

callbacks are only supported in spmd computations when all mesh axes are partitioned manually (no partial automatic sharding).

What it means

When lowering under SPMD (jax.jit with a Mesh / sharding), callbacks can only be executed with per-device semantics, which requires every mesh axis to be in manual_axes (fully manual sharding). If only some axes are manual while others are automatically partitioned, there is no consistent way to place and run the callback, so _callback_op_sharding raises NotImplementedError.

Source

Thrown at jax/_src/callback.py:158

  ndims = [0]
  if avals:
    ndims = [x.ndim for x in avals if isinstance(x, core.ShapedArray)]
  return SdyArrayList(tuple(
      SdyArray(
          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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all mesh axes manual around the callback: run it inside shard_map (jax.experimental.shard_map.shard_map), which sets manual_axes to the full mesh
2. Move the callback outside the sharded jit region (compute on host with device_get)
3. Restructure so the mesh over which the callback executes is fully manual (split into per-stage meshes)

Example fix

# before
@jax.jit
def f(x):  # x sharded over mesh with auto axes
    return pure_callback(fn, x, result_dtype=x.dtype)  # NotImplementedError

# after
from jax.experimental.shard_map import shard_map
@jax.jit
def f(x):
    return shard_map(lambda xb: pure_callback(fn, xb, result_dtype=xb.dtype),
                     mesh, in_specs=(P('data',), out_specs=P('data',))(x)

from jax._src import sharding_impls
ctx_axes = jax.experimental.multiaddr_utils  # in practice: check your mesh usage
# practical pre-check inside manual code:
assert manual_axes == frozenset(mesh.axis_names), (
    'callbacks require fully manual mesh; wrap in shard_map')

try:
    jax.jit(f)(x)
except NotImplementedError as e:
    if 'all mesh axes are partitioned manually' in str(e):
        f = shard_map(f, mesh, in_specs=..., out_specs=...)
    raise

When it happens

Trigger: Using jax.pure_callback or jax.io_callback inside a jax.jit computation with a Mesh where manual_axes (from autodiff/manual partitioning, e.g. jax.lax.map with in_axes, pjit with partially manual axes, or shard_map boundaries) is a strict subset of mesh.axis_names.

Common situations: Combining shard_map-like manual regions with automatic GSPMD sharding in the same computation and calling a callback in the auto-sharded part; migrating pmap/pjit code with callbacks to newer SPMD APIs; using callbacks inside jax.debug with mixed manual/auto axes.

Related errors

• callbacks do not support specifying sharding inside spmd com
• pallas_call requires all mesh axes to be Manual, got {get_ab
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• pure_callback only supports SingleDeviceSharding, but got {t

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