Fix test_make_array_from_single_device_arrays in random_test.py for TPU 7x. This is because create_mesh return a different order of devices while arrays are being created with jax.devices() order.

PiperOrigin-RevId: 761975228

Author: yashk2810

tests/random_test.py

@@ -991,7 +991,7 @@ def callback(index):
   def test_make_array_from_single_device_arrays(self):
     devices = jax.devices()
     shape = (len(devices),)
-    mesh = jtu.create_mesh((len(devices),), ('x',))
+    mesh = jtu.create_mesh((len(devices),), ('x',), iota_order=True)
     sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec('x'))
     keys = random.split(random.key(0), len(devices))
     arrays = [jax.device_put(keys[i:i + 1], device) for i, device in enumerate(devices)]