Consolidate the distributed fused attention tests to shared input gen…

…eration and execition logic.

Signed-off-by: Michael Goldfarb <[email protected]>

tests/jax/distributed_test_base.py

@@ -17,37 +17,30 @@
 
 def generate_configs():
     configs = []
+    mr = MeshResource(dp_resource="dp", tp_resource="tp")
+    axes = ("dp", "tp")
     if is_devices_enough(2):
-        configs.append([2, (2,), "dp", MeshResource(dp_resource="dp")])
-        configs.append([2, (2,), "tp", MeshResource(tp_resource="tp")])
+        configs.append(pytest.param(2, (2, 1), axes, mr, id="n2_dp2_tp1"))
+        configs.append(pytest.param(2, (1, 2), axes, mr, id="n2_dp1_tp2"))
 
     if is_devices_enough(4):
-        TP_size = 2
-        DP_size = 2
-        configs.append(
-            [4, (DP_size, TP_size), ("dp", "tp"), MeshResource(dp_resource="dp", tp_resource="tp")]
-        )
+        configs.append(pytest.param(4, (2, 2), axes, mr, id=f"n4_dp2_tp2"))
 
     return configs
 
 
 def generate_context_parallel_configs():
     configs = []
-
+    mr = MeshResource(dp_resource="dp", cp_resource="cp", tp_resource="tp")
+    axes = ("dp", "cp", "tp")
     DP_sizes = (1, 2)
     CP_sizes = (1, 2, 4, 8)
     TP_sizes = (1, 2)
     for dp, cp, tp in product(DP_sizes, CP_sizes, TP_sizes):
         ndev = cp * tp * dp
         if is_devices_enough(ndev):
             configs.append(
-                pytest.param(
-                    ndev,
-                    (dp, cp, tp),
-                    ("dp", "cp", "tp"),
-                    MeshResource(dp_resource="dp", cp_resource="cp", tp_resource="tp"),
-                    id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}",
-                )
+                pytest.param(ndev, (dp, cp, tp), axes, mr, id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}")
             )
 
     return configs