Sort devices by their implicit order instead of explicitly by id. IDs may be randomly generated, so it's better to rely on the implicit order, which is currently based on (process index, id).

PiperOrigin-RevId: 650294623

Author: T5X Team

t5x/test_utils.py

@@ -233,7 +233,7 @@ def create_global_mesh(mesh_shape, axis_names):
   size = np.prod(mesh_shape)
   if len(jax.devices()) < size:
     raise unittest.SkipTest(f'Test requires {size} global devices.')
-  devices = sorted(jax.devices(), key=lambda d: d.id)
+  devices = sorted(jax.devices())
   mesh_devices = np.array(devices[:size]).reshape(mesh_shape)
   global_mesh = Mesh(mesh_devices, axis_names)
   return global_mesh