[JAX] Extend colocated_cpu_devices to accept Mesh besides devices

jax/experimental/colocated_python/api.py

@@ -16,25 +16,43 @@
 from __future__ import annotations
 
 import collections
-from typing import Any, Callable, Sequence, Type
+from typing import Any, Callable, Sequence, Type, overload
 
 import jax
 from jax._src import api_util
 from jax._src import util
 from jax.experimental.colocated_python.func import make_callable
 from jax.experimental.colocated_python.obj import wrap_class
+import numpy as np
 
 
+@overload
 def colocated_cpu_devices(
-    devices: Sequence[jax.Device],
+    devices_or_mesh: Sequence[jax.Device],
 ) -> Sequence[jax.Device]:
-  """Finds CPU devices colocated with the given devices."""
-  if not isinstance(devices, tuple):
-    devices = tuple(devices)
+  ...
+
+
+@overload
+def colocated_cpu_devices(
+    devices_or_mesh: jax.sharding.Mesh,
+) -> jax.sharding.Mesh:
+  ...
+
+
+def colocated_cpu_devices(devices_or_mesh):
+  """Finds devices or a mesh that has CPU devices colocated with the given devices or mesh."""
+  if isinstance(devices_or_mesh, jax.sharding.Mesh):
+    return _colocated_cpu_mesh_cached(devices_or_mesh)
+
+  if not isinstance(devices_or_mesh, tuple):
+    devices_or_mesh = tuple(devices_or_mesh)
   try:
-    return _colocated_cpu_devices_cached(devices)
+    return _colocated_cpu_devices_cached(devices_or_mesh)
   except (ValueError, AttributeError):
-    return _colocated_cpu_devices_cached_fallback_to_cpu_backend(devices)
+    return _colocated_cpu_devices_cached_fallback_to_cpu_backend(
+        devices_or_mesh
+    )
 
 
 @util.cache(max_size=1024, trace_context_in_key=False)
@@ -78,6 +96,20 @@ def _colocated_cpu_devices_cached_fallback_to_cpu_backend(
   ]
 
 
+@util.cache(max_size=1024, trace_context_in_key=False)
+def _colocated_cpu_mesh_cached(mesh: jax.sharding.Mesh) -> jax.sharding.Mesh:
+  """Returns a CPU mesh that is similar to the given mesh but has colocated CPU devices."""
+  # Finding colocated CPU devices reuses the cache of `colocated_cpu_devices`
+  # called with devices. `_colocated_cpu_mesh` itself is also cached to avoid
+  # creating a new `Mesh` object repeatedly.
+  flat_cpu_devices = colocated_cpu_devices(tuple(mesh.devices.flat))
+  return jax.sharding.Mesh(
+      np.array(flat_cpu_devices).reshape(mesh.axis_sizes),
+      mesh.axis_names,
+      axis_types=mesh.axis_types,
+  )
+
+
 def colocated_python(fun: Callable[..., Any]) -> Callable[..., Any]:
   """Executes the given Python function on the same devices as the arguments."""
   return make_callable(

tests/colocated_python_test.py

@@ -58,6 +58,20 @@ def setUp(self):
         " requires NumPy 2.0.0 or later"
       )
 
+  def testColocatedCpuDevices(self):
+    mesh = jax.sharding.Mesh(
+        np.array(jax.local_devices()[:1]).reshape((1, 1)), ("x", "y")
+    )
+    cpu_mesh1 = colocated_python.colocated_cpu_devices(mesh)
+
+    cpu_devices = colocated_python.colocated_cpu_devices(
+        jax.local_devices()[:1]
+    )
+    cpu_mesh2 = jax.sharding.Mesh(
+        np.array(cpu_devices).reshape((1, 1)), ("x", "y")
+    )
+    self.assertEqual(cpu_mesh1, cpu_mesh2)
+
   def testMakeColocatedPythonProgram(self):
     def add_one(x):
       return x + 1