Merge pull request #1671 from AI-Hypercomputer:mattdavidow-remove-shmap-test

PiperOrigin-RevId: 754072676

Author: maxtext authors

MaxText/tests/integration_tests/shmap_collective_matmul_test.py

@@ -22,12 +22,15 @@
 from MaxText.globals import PKG_DIR
 
 sys.path.append(os.path.join(os.path.dirname(PKG_DIR), "pedagogical_examples"))
-from pedagogical_examples.shmap_collective_matmul import main
 
+# Uncomment the import when b/415022795 is fixed
+#from pedagogical_examples.shmap_collective_matmul import main
 
+
+@pytest.mark.skip(reason="Enable when b/415022795 is fixed")
 @pytest.mark.integration_test
 @pytest.mark.tpu_only
 def test_shmap_collective_matmul_example():
   """Validate Pedagogical Example, Shmap_collective_matmul."""
-
-  assert main() is True
+  # Uncomment main() assertion when b/415022795 is fixed
+  #assert main() is True

MaxText/tests/kernels_test.py

@@ -35,14 +35,15 @@ class RaggedAttentionTest(unittest.TestCase):
   head_dim = 128
 
   dtype = jnp.float32
-  key = jax.random.key(0)
-  k1, k2, k3 = jax.random.split(key, 3)
 
   @pytest.mark.tpu_only
   def test_ragged_mqa(self):
-    q = jax.random.normal(self.k1, (self.batch_size, 1, self.head_dim), dtype=self.dtype)
-    k = jax.random.normal(self.k2, (self.batch_size, self.max_target_length, self.head_dim), dtype=self.dtype)
-    v = jax.random.normal(self.k3, (self.batch_size, self.max_target_length, self.head_dim), dtype=self.dtype)
+    key = jax.random.key(0)
+    k1, k2, k3 = jax.random.split(key, 3)
+
+    q = jax.random.normal(k1, (self.batch_size, 1, self.head_dim), dtype=self.dtype)
+    k = jax.random.normal(k2, (self.batch_size, self.max_target_length, self.head_dim), dtype=self.dtype)
+    v = jax.random.normal(k3, (self.batch_size, self.max_target_length, self.head_dim), dtype=self.dtype)
     lengths = jnp.array(np.random.randint(1, self.max_target_length, self.batch_size), dtype=jnp.int32)
 
     ragged_out, ragged_max, ragged_denom = ragged_mqa(q, k, v, lengths)
@@ -58,12 +59,15 @@ def test_ragged_mqa(self):
 
   @pytest.mark.tpu_only
   def test_ragged_mha(self):
-    q = jax.random.normal(self.k1, (self.batch_size, 1, self.num_query_heads, self.head_dim), dtype=self.dtype)
+    key = jax.random.key(0)
+    k1, k2, k3 = jax.random.split(key, 3)
+
+    q = jax.random.normal(k1, (self.batch_size, 1, self.num_query_heads, self.head_dim), dtype=self.dtype)
     k = jax.random.normal(
-        self.k2, (self.batch_size, self.max_target_length, self.num_query_heads, self.head_dim), dtype=self.dtype
+        k2, (self.batch_size, self.max_target_length, self.num_query_heads, self.head_dim), dtype=self.dtype
     )
     v = jax.random.normal(
-        self.k3, (self.batch_size, self.max_target_length, self.num_query_heads, self.head_dim), dtype=self.dtype
+        k3, (self.batch_size, self.max_target_length, self.num_query_heads, self.head_dim), dtype=self.dtype
     )
     lengths = jnp.array(np.random.randint(1, self.max_target_length, self.batch_size), dtype=jnp.int32)
 
@@ -81,12 +85,15 @@ def test_ragged_mha(self):
 
   @pytest.mark.tpu_only
   def test_ragged_gqa(self):
-    q = jax.random.normal(self.k1, (self.batch_size, 1, self.num_query_heads, self.head_dim), dtype=self.dtype)
+    key = jax.random.key(0)
+    k1, k2, k3 = jax.random.split(key, 3)
+
+    q = jax.random.normal(k1, (self.batch_size, 1, self.num_query_heads, self.head_dim), dtype=self.dtype)
     k = jax.random.normal(
-        self.k2, (self.batch_size, self.max_target_length, self.num_kv_heads, self.head_dim), dtype=self.dtype
+        k2, (self.batch_size, self.max_target_length, self.num_kv_heads, self.head_dim), dtype=self.dtype
     )
     v = jax.random.normal(
-        self.k3, (self.batch_size, self.max_target_length, self.num_kv_heads, self.head_dim), dtype=self.dtype
+        k3, (self.batch_size, self.max_target_length, self.num_kv_heads, self.head_dim), dtype=self.dtype
     )
     lengths = jnp.array(np.random.randint(1, self.max_target_length, self.batch_size), dtype=jnp.int32)
 

pedagogical_examples/shmap_collective_matmul.py

@@ -35,6 +35,7 @@
 MESH_FSDP_AXIS = "fsdp"
 MESH_TENSOR_AXIS = "tp"
 
+# We should not call jax.devices() when this file is imported b/415022795.
 d = jax.devices()
 outd = [[[d[0], d[1], d[3], d[2]]]]
 global_mesh = Mesh(outd, (MESH_DATA_AXIS, MESH_FSDP_AXIS, MESH_TENSOR_AXIS))