jax-ml/jax · error · NotImplementedError

callbacks are only supported in spmd computations when all…

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.

Solutions

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
Move the callback outside the sharded jit region (compute on host with device_get)
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)

Defensive patterns

Strategy: validation

Validate before calling

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 / catch

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

Prevention

Always run callbacks inside shard_map when using SPMD meshes
Avoid mixing manual and automatic axes in one jit scope
Test SPMD programs with callbacks on small meshes first

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

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

Appendix: 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)