[JAX] Fix colocated_python.colocated_cpu_devices() for multi-host CPU setups

On McJAX, `colocated_python.colocated_cpu_devices()` was querying only local CPU devices when finding colocated devices for (global) TPU devices. The API should have queried all CPU devices regardless of their addressability.

PiperOrigin-RevId: 833889551

Author: hyeontaek

jax/experimental/colocated_python/api.py

@@ -115,7 +115,7 @@ def _colocated_cpu_devices_cached_fallback_to_cpu_backend(
   else:
     # PjRt-IFRT on a non-CPU platform currently defines CPU devices on a separae
     # CPU backend.
-    cpu_backend_devices = jax.local_devices(backend="cpu")
+    cpu_backend_devices = jax.devices(backend="cpu")
   device_index_map = {device.id: i for i, device in enumerate(jax.devices())}
 
   available_devices = devices[: min(len(cpu_backend_devices), len(devices))]

tests/multiprocess/BUILD

@@ -78,6 +78,25 @@ jax_multiprocess_test(
     ],
 )
 
+jax_multiprocess_test(
+    name = "colocated_python_test",
+    srcs = ["colocated_python_test.py"],
+    disable_configs = [
+        # This config has two cores per chip, and JAX distributed does not get
+        # the correct number of logical devices per host.
+        "tpu_v3_x4",
+    ],
+    enable_backends = [
+        "gpu",
+        "tpu",
+    ],
+    main = "colocated_python_test.py",
+    deps = [
+        "//jax:experimental_colocated_python",
+        "//jax/_src:test_multiprocess",
+    ],
+)
+
 jax_multiprocess_test(
     name = "device_id_test",
     srcs = ["device_id_test.py"],

tests/multiprocess/colocated_python_test.py

@@ -0,0 +1,65 @@
+# Copyright 2025 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Multihost tests for jax.Array."""
+
+import jax
+from jax._src import test_multiprocess as jt_multiprocess
+from jax._src import test_util as jtu
+from jax.experimental import colocated_python
+import numpy as np
+
+
+class ColocatedPythonTestMultiHost(jt_multiprocess.MultiProcessTest):
+
+  def setUp(self):
+    super().setUp()
+    jtu.request_cpu_devices(jax.local_device_count())
+
+  def test_colocated_cpu_devices(self):
+    if jax.device_count() % 2 == 0:
+      mesh_shape = (2, jax.device_count() // 2)
+    else:
+      mesh_shape = (1, jax.device_count())
+    mesh = jax.make_mesh(mesh_shape, ("x", "y"))
+    cpu_mesh1 = colocated_python.colocated_cpu_devices(mesh)
+
+    cpu_devices = colocated_python.colocated_cpu_devices(mesh.devices.flat)
+    cpu_mesh2 = jax.make_mesh(mesh_shape, ("x", "y"), devices=cpu_devices)
+    self.assertEqual(cpu_mesh1, cpu_mesh2)
+
+  def test_simple_function(self):
+    @colocated_python.colocated_python
+    def add_one(x):
+      return jax.make_array_from_single_device_arrays(
+          x.shape, x.sharding, [s.data + 1 for s in x.addressable_shards])
+
+    mesh = jax.make_mesh((jax.device_count(),), ("x",))
+    cpu_mesh = colocated_python.colocated_cpu_devices(mesh)
+    cpu_sharding = jax.NamedSharding(cpu_mesh, jax.P("x"))
+
+    x = np.arange(cpu_mesh.size)
+    x = jax.device_put(x, cpu_sharding)
+
+    out = add_one(x)
+
+    out = jax.jit(lambda x: x,
+                  out_shardings=jax.NamedSharding(cpu_mesh, jax.P()))(out)
+    out = jax.device_get(out)
+
+    np.testing.assert_equal(out, np.arange(cpu_mesh.size) + 1)
+
+
+if __name__ == "__main__":
+  jt_multiprocess.main()