{"record":{"id":"b2449c3ee3bab44c","repo":"jax-ml/jax","slug":"pure-callback-only-supports-singledevicesharding","errorCode":null,"errorMessage":"pure_callback only supports SingleDeviceSharding, but got {type(sharding)}","messagePattern":"pure_callback only supports SingleDeviceSharding, but got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/callback.py","lineNumber":180,"sourceCode":"    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\"\n            f\" {type(sharding)}\"\n        )\n      device = next(iter(sharding.device_set))\n      device_assignment = axis_context.device_assignment\n      if device_assignment is None:\n        raise AssertionError(\n            \"Please file a bug at https://github.com/jax-ml/jax/issues\")\n      try:\n        device_index = device_assignment.index(device)\n      except IndexError as e:\n        raise ValueError(\n            \"Sharding provided to pure_callback specifies a device\"\n            f\" {device} that is not in the device assignment\"\n            f\" ({device_assignment})\") from e\n    else:\n      device_index = 0\n","sourceCodeStart":162,"sourceCodeEnd":198,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/callback.py#L162-L198","documentation":"Outside SPMD (plain ShardingContext lowering), pure_callback lowers via an OpSharding that places computation on a single device; therefore only SingleDeviceSharding (or a scalar NamedSharding, which degenerates to one device) is accepted. Any other sharding type — NamedSharding over a real mesh, PositionalSharding, MultiDeviceSharding — hits NotImplementedError with the offending type name.","triggerScenarios":"Passing sharding=NamedSharding(mesh, P('data','model')), PositionalSharding, GSPMDSharding, etc. to jax.pure_callback or jax.io_callback in a normal (non-SPMD) jax.jit context.","commonSituations":"Assuming pure_callback can express multi-device output sharding; reusing a pjit out_shards object as the callback sharding; gradually sharding a pipeline and passing the mesh sharding to a host callback.","solutions":["Pass sharding=None (default; JAX infers a single device) when multi-device sharding is not truly needed","Use SingleDeviceSharding for an explicit device: sharding=jax.sharding.SingleDeviceSharding(jax.devices()[0])","For multi-device outputs, run the callback per-shard via shard_map (manual SPMD) without an explicit sharding argument"],"exampleFix":"# before\nsh = NamedSharding(mesh, P('data'))\nout = pure_callback(fn, x, result_dtype=x.dtype, sharding=sh)  # NotImplementedError\n\n# after\nfrom jax.sharding import SingleDeviceSharding\nout = pure_callback(fn, x, result_dtype=x.dtype,\n                    sharding=SingleDeviceSharding(jax.devices()[0]))","handlingStrategy":"type-guard","validationCode":"from jax.sharding import SingleDeviceSharding, NamedSharding\n\ndef ok_sharding(s):\n    if s is None: return True\n    if isinstance(s, SingleDeviceSharding): return True\n    if isinstance(s, NamedSharding) and s.mesh.is_scalar: return True\n    return False\nassert ok_sharding(sh), f'pure_callback sharding must be single-device, got {type(sharding)}'","typeGuard":"from jax.sharding import SingleDeviceSharding, NamedSharding\ndef is_pure_callback_sharding(s) -> bool:\n    return s is None or isinstance(s, SingleDeviceSharding) or (\n        isinstance(s, NamedSharding) and bool(s.mesh.is_scalar))","tryCatchPattern":"try:\n    pure_callback(fn, x, result_dtype=d, sharding=sh)\nexcept NotImplementedError as e:\n    if 'only supports SingleDeviceSharding' in str(e):\n        sh = SingleDeviceSharding(next(iter(sh.device_set)))\n        pure_callback(fn, x, result_dtype=d, sharding=sh)\n    raise","preventionTips":["Default to sharding=None for pure_callback","Never reuse pjit output shardings as callback shardings","Switch to shard_map for genuinely multi-device callback outputs"],"tags":["jax","sharding","pure-callback","single-device-sharding"],"backgroundTag":"unsupported-sharding-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}