Sort devices explicitly by process index, then id (as opposed to IDs alone). IDs may be randomly generated, and are not guaranteed to be ordered based on their process index.

PiperOrigin-RevId: 650294623

Author: ChexDev

chex/_src/asserts_test.py

@@ -1200,12 +1200,14 @@ def test_assert_tree_is_on_device(self):
       asserts.assert_tree_is_on_device(tpu_2_tree, device=tpu_2)
 
       with self.assertRaisesRegex(
-          AssertionError, _get_err_regex(r"'a' resides on.*TpuDevice\(id=0")
+          AssertionError,
+          _get_err_regex(r"'a' resides on.*TpuDevice\(process_index=0, id=0"),
       ):
         asserts.assert_tree_is_on_device(tpu_1_tree, device=tpu_2)
 
       with self.assertRaisesRegex(
-          AssertionError, _get_err_regex(r"'a' resides on.*TpuDevice\(id=1")
+          AssertionError,
+          _get_err_regex(r"'a' resides on.*TpuDevice\(process_index=0, id=1"),
       ):
         asserts.assert_tree_is_on_device(tpu_2_tree, device=tpu_1)
 
@@ -1735,7 +1737,7 @@ def test_assert_equal_pass(self, first, second):
     asserts.assert_equal(first, second)
 
   def test_assert_equal_pass_on_arrays(self):
-    # Not using named_parameters, becase JAX cannot be used before app.run().
+    # Not using named_parameters, because JAX cannot be used before app.run().
     asserts.assert_equal(jnp.ones([]), np.ones([]))
     asserts.assert_equal(
         jnp.ones([], dtype=jnp.int32), np.ones([], dtype=np.float64))