{"record":{"id":"f87282341383122e","repo":"jax-ml/jax","slug":"callbacks-do-not-support-specifying-sharding-insid","errorCode":null,"errorMessage":"callbacks do not support specifying sharding inside spmd computations","messagePattern":"callbacks do not support specifying sharding inside spmd computations","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":163,"sourceCode":"          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\n          sharding.mesh.is_scalar):  # pyrefly: ignore[missing-attribute]\n        pass\n      elif not isinstance(sharding, SingleDeviceSharding):\n        raise NotImplementedError(\n            \"pure_callback only supports SingleDeviceSharding, but got\"","sourceCodeStart":145,"sourceCodeEnd":181,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L145-L181","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","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"],"exampleFix":"# before\n@shard_map(mesh=mesh, in_specs=P('d'), out_specs=P('d'))\ndef f(x):\n    return pure_callback(fn, x, result_dtype=x.dtype,\n                         sharding=NamedSharding(mesh, P('d')))  # raises\n\n# after\n@shard_map(mesh=mesh, in_specs=P('d'), out_specs=P('d'))\ndef f(x):\n    return pure_callback(fn, x, result_dtype=x.dtype)  # no sharding arg","handlingStrategy":"validation","validationCode":"def spmd_callback(fn, x, sharding):\n    if in_manual_spmd():  # e.g. flag set inside shard_map body\n        assert sharding is None, 'do not pass sharding under manual SPMD'\n    return pure_callback(fn, x, result_dtype=x.dtype, sharding=sharding)","typeGuard":null,"tryCatchPattern":"try:\n    jax.jit(f)(x)\nexcept NotImplementedError as e:\n    if 'do not support specifying sharding inside spmd' in str(e):\n        retry with sharding=None\n    raise","preventionTips":["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"],"tags":["jax","spmd","sharding","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"}