{"record":{"id":"ef188add32f79778","repo":"jax-ml/jax","slug":"callbacks-are-only-supported-in-spmd-computations","errorCode":null,"errorMessage":"callbacks are only supported in spmd computations when all mesh axes are partitioned manually (no partial automatic sharding).","messagePattern":"callbacks are only supported in spmd computations when all mesh axes are partitioned manually \\(no partial automatic sharding\\)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":158,"sourceCode":"  ndims = [0]\n  if avals:\n    ndims = [x.ndim for x in avals if isinstance(x, core.ShapedArray)]\n  return SdyArrayList(tuple(\n      SdyArray(\n          mesh_shape=(),\n          dim_shardings=(SdyDim(axes=(), is_open=False),) * ndim,\n          logical_device_ids=())\n      for ndim in ndims))\n\n\ndef _callback_op_sharding(\n    axis_context, sharding: Sharding | None, avals_out\n):\n  if isinstance(axis_context, sharding_impls.SPMDAxisContext):\n    # If we have fully manual sharding during lowering, that means the JAX\n    # program has per-device semantics, so we run the callback on each device.\n    if axis_context.manual_axes != frozenset(axis_context.mesh.axis_names):\n      raise NotImplementedError(\n          \"callbacks are only supported in spmd computations when all mesh\"\n          \" axes are partitioned manually (no partial automatic sharding).\"\n      )\n    if sharding is not None:\n      raise NotImplementedError(\n          \"callbacks do not support specifying sharding inside spmd\"\n          \" computations\"\n      )\n    if config.use_shardy_partitioner.value:\n      op_sharding = _get_sdy_array_list_for_callbacks(avals_out)\n    else:\n      op_sharding = xc.OpSharding()\n      op_sharding.type = xc.OpSharding.Type.MANUAL\n    return op_sharding\n\n  if isinstance(axis_context, sharding_impls.ShardingContext):\n    if sharding is not None:\n      if (isinstance(sharding, sharding_impls.NamedSharding) and","sourceCodeStart":140,"sourceCodeEnd":176,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L140-L176","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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)"],"exampleFix":"# before\n@jax.jit\ndef f(x):  # x sharded over mesh with auto axes\n    return pure_callback(fn, x, result_dtype=x.dtype)  # NotImplementedError\n\n# after\nfrom jax.experimental.shard_map import shard_map\n@jax.jit\ndef f(x):\n    return shard_map(lambda xb: pure_callback(fn, xb, result_dtype=xb.dtype),\n                     mesh, in_specs=(P('data',), out_specs=P('data',))(x)","handlingStrategy":"validation","validationCode":"from jax._src import sharding_impls\nctx_axes = jax.experimental.multiaddr_utils  # in practice: check your mesh usage\n# practical pre-check inside manual code:\nassert manual_axes == frozenset(mesh.axis_names), (\n    'callbacks require fully manual mesh; wrap in shard_map')","typeGuard":null,"tryCatchPattern":"try:\n    jax.jit(f)(x)\nexcept NotImplementedError as e:\n    if 'all mesh axes are partitioned manually' in str(e):\n        f = shard_map(f, mesh, in_specs=..., out_specs=...)\n    raise","preventionTips":["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"],"tags":["jax","spmd","sharding","mesh","pure-callback","manual-axes"],"backgroundTag":"jax-spmd-manual-sharding-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}