{"record":{"id":"6447e0882570867b","repo":"jax-ml/jax","slug":"for-a-cross-host-reshard-in-multi-controller-jax","errorCode":null,"errorMessage":"For a cross-host reshard in multi-controller JAX, input and target sharding should have the same set of devices. Got input's device set ids: {inp_ids} on platform {inp_plat} and target sharding's device set ids: {target_ids} on platform {target_plat}.\n\nThere is experimental support for cross-host transfers with different device sets, when input/output shardings have the same indices and layouts, in the TFRT TPU runtime only.","messagePattern":"For a cross-host reshard in multi-controller JAX, input and target sharding should have the same set of devices\\. Got input's device set ids: (.+?) on platform (.+?) and target sharding's device set ids: (.+?) on platform (.+?)\\.\n\nThere is experimental support for cross-host transfers with different device sets, when input/output shardings have the same indices and layouts, in the TFRT TPU runtime only\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/dispatch.py","lineNumber":497,"sourceCode":"        s.device_set == x_sharding.device_set):\n      assert isinstance(s, NamedSharding), s\n      return _different_device_order_reshard(x, s, copy)\n\n    if (x_is_jax_array and x._committed and xb.process_count() > 1\n        and _is_supported_cross_host_transfer(x.ndim, x_sharding, s)):\n      return _DeferredCrossHostTransferArg(x, s, copy)\n\n    if not s_is_fully_addressable:\n      # If both the source and target shardings are not fully addressable and\n      # one of the above conditions has not been met, then assume that the user\n      # is attempting a different device order reshard.\n      if (x_is_jax_array and not x_is_fully_addressable\n          and s.device_set != x_sharding.device_set):\n        inp_ids = [d.id for d in x_sharding._device_assignment]\n        inp_plat = x_sharding._device_assignment[0].platform.upper()\n        target_ids = [d.id for d in s._device_assignment]\n        target_plat = s._device_assignment[0].platform.upper()\n        raise ValueError(\n            \"For a cross-host reshard in multi-controller JAX, input and target\"\n            \" sharding should have the same set of devices. Got input's device\"\n            f\" set ids: {inp_ids} on platform {inp_plat} and target sharding's\"\n            f\" device set ids: {target_ids} on platform {target_plat}.\\n\\n\"\n            \"There is experimental support for cross-host transfers with \"\n            \"different device sets, when input/output shardings have the same \"\n            \"indices and layouts, in the TFRT TPU runtime only.\")\n\n      if ((x_is_jax_array and not x._committed) or\n          type(x) in array_types or type(x) in dtypes.python_scalar_types):\n        # If all hosts participate in the sharding, assert that the input is the\n        # same on all hosts. If some hosts have no addressable devices in the\n        # sharding, bypass the check, since we can't easily distinguish between\n        # these two cases: (1) the sharding contains the same subset of global\n        # devices on all hosts (and hosts with no addressable devices in the\n        # sharding do not transfer data) or (2) the sharding contains a\n        # different subset of devices on each host. For (1), the input should be\n        # the same on all hosts, but for (2) it need not be.","sourceCodeStart":479,"sourceCodeEnd":515,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/dispatch.py#L479-L515","documentation":"In multi-controller (multi-process) JAX, resharding an array across hosts requires the input array's device set and the target Sharding's device set to match. This error reports the mismatched device ids/platforms and notes the only experimental exception (TFRT TPU runtime with identical indices/layouts).","triggerScenarios":"jax.device_put where x is a non-fully-addressable jax.Array whose sharding's device_set differs from the target sharding's device_set — e.g. sharding built over a different mesh/device order than the array's current placement.","commonSituations":"Constructing NamedSharding/GSPMDSharding from devices in a different order than the mesh the array was created on; mixing single-process and multi-process device assignments; using a global device list that differs between input and target in a multi-host job.","solutions":["Rebuild the target sharding from the same device set/order as the array's current sharding (e.g. via mesh.devices or x.sharding.device_set)","Ensure jax.distributed.initialize uses identical num_processes/process_id ordering so device assignments agree across hosts","If on TFRT TPU runtime, ensure input/output shardings have identical indices and layouts to use the experimental path"],"exampleFix":"# before\ntarget = jax.sharding.NamedSharding(jax.devices()[::-1], P('x'))  # different order\ny = jax.device_put(x, target)\n# after\ntarget = jax.sharding.NamedSharding(jax.devices(), P('x'))  # same set/order as x\ny = jax.device_put(x, target)","handlingStrategy":"validation","validationCode":"def same_device_set(a_sharding, target):\n    return a_sharding.device_set == target.device_set\nassert same_device_set(x.sharding, target_sharding), 'device sets differ'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build input and target shardings from the same mesh","Keep device ordering consistent across processes","Log device_set ids before reshards in multi-host jobs"],"tags":["jax","multi-controller","resharding","device-mismatch"],"backgroundTag":"distributed-sharding-misconfiguration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}