jax-ml/jax · error · ValueError

For a cross-host reshard in multi-controller JAX, input and

Error message

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}.

There 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.

What it means

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).

Source

Thrown at jax/_src/dispatch.py:497

        s.device_set == x_sharding.device_set):
      assert isinstance(s, NamedSharding), s
      return _different_device_order_reshard(x, s, copy)

    if (x_is_jax_array and x._committed and xb.process_count() > 1
        and _is_supported_cross_host_transfer(x.ndim, x_sharding, s)):
      return _DeferredCrossHostTransferArg(x, s, copy)

    if not s_is_fully_addressable:
      # If both the source and target shardings are not fully addressable and
      # one of the above conditions has not been met, then assume that the user
      # is attempting a different device order reshard.
      if (x_is_jax_array and not x_is_fully_addressable
          and s.device_set != x_sharding.device_set):
        inp_ids = [d.id for d in x_sharding._device_assignment]
        inp_plat = x_sharding._device_assignment[0].platform.upper()
        target_ids = [d.id for d in s._device_assignment]
        target_plat = s._device_assignment[0].platform.upper()
        raise ValueError(
            "For a cross-host reshard in multi-controller JAX, input and target"
            " sharding should have the same set of devices. Got input's device"
            f" set ids: {inp_ids} on platform {inp_plat} and target sharding's"
            f" device set ids: {target_ids} on platform {target_plat}.\n\n"
            "There 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.")

      if ((x_is_jax_array and not x._committed) or
          type(x) in array_types or type(x) in dtypes.python_scalar_types):
        # If all hosts participate in the sharding, assert that the input is the
        # same on all hosts. If some hosts have no addressable devices in the
        # sharding, bypass the check, since we can't easily distinguish between
        # these two cases: (1) the sharding contains the same subset of global
        # devices on all hosts (and hosts with no addressable devices in the
        # sharding do not transfer data) or (2) the sharding contains a
        # different subset of devices on each host. For (1), the input should be
        # the same on all hosts, but for (2) it need not be.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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)
2. Ensure jax.distributed.initialize uses identical num_processes/process_id ordering so device assignments agree across hosts
3. If on TFRT TPU runtime, ensure input/output shardings have identical indices and layouts to use the experimental path

Example fix

# before
target = jax.sharding.NamedSharding(jax.devices()[::-1], P('x'))  # different order
y = jax.device_put(x, target)
# after
target = jax.sharding.NamedSharding(jax.devices(), P('x'))  # same set/order as x
y = jax.device_put(x, target)

def same_device_set(a_sharding, target):
    return a_sharding.device_set == target.device_set
assert same_device_set(x.sharding, target_sharding), 'device sets differ'

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Specified {device=} which requires a copy since the source d
• Specified {device=} which requires a copy since the source d
• top_level_all_gather doesn't allow input {aval} to be unshar
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64

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