fix: L0 placeholder shapes are invalid empty tuple/list

[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_distributed_fused_attn.py: L0 placeholder shapes are invalid empty tuple/list.

Changes

• tests/jax/test_distributed_fused_attn.py: L0 placeholder shapes are invalid empty tuple/list.

Details

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -1,21 +1,21 @@
-DISTRIBUTED_SELF_ATTN_DATA_SHAPES = {
-    "L0": [()],
-    "L1": [(32, 1024, 16, 128)],
-    "L2": [(32, 512, 12, 64)],
-}
-
-...
-
-DISTRIBUTED_CROSS_ATTN_DATA_SHAPES = {
-    "L0": [()],
-    "L1": [[32, 512, 16, 64]],
-    "L2": [[32, 128, 12, 64]],
-}
-
-...
-
-REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = {
-    "L0": [[]],
-    "L1": [[3, 32, 8, 64]],
-    "L2": [[4, 32, 12, 32], [1, 16, 1, 1]],
-}
+DISTRIBUTED_SELF_ATTN_DATA_SHAPES = {
+    "L0": [],
+    "L1": [(32, 1024, 16, 128)],
+    "L2": [(32, 512, 12, 64)],
+}
+
+...
+
+DISTRIBUTED_CROSS_ATTN_DATA_SHAPES = {
+    "L0": [],
+    "L1": [[32, 512, 16, 64]],
+    "L2": [[32, 128, 12, 64]],
+}
+
+...
+
+REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = {
+    "L0": [],
+    "L1": [[3, 32, 8, 64]],
+    "L2": [[4, 32, 12, 32], [1, 16, 1, 1]],
+}

Tests

• tests/jax/test_distributed_fused_attn.py

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -437,3 +437,11 @@ class TestReorderCausalLoadBalancing:
 
         reordered = reorder(tensor, reorder_strategy, cp_size, seq_dim, stripe_size)
         inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim, stripe_size)
 
         assert jnp.array_equal(inversed, ref)
+
+
+def test_l0_shape_placeholders_are_empty():
+    """Ensure L0 placeholders are empty collections and not invalid empty shapes."""
+    assert DISTRIBUTED_SELF_ATTN_DATA_SHAPES.get("L0") == []
+    assert DISTRIBUTED_CROSS_ATTN_DATA_SHAPES.get("L0") == []
+    assert REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES.get("L0") == []

[greptile-apps]

Greptile Summary

This PR replaces invalid L0 placeholder shapes with empty parameter collections and adds a regression test enforcing that representation.

• Updates the self-attention, cross-attention, and causal load-balancing L0 shape sets.
• Preserves the nonempty L1 and L2 distributed test cases.
• Adds assertions covering all three L0 placeholders.

Confidence Score: 5/5

The PR appears safe to merge with no actionable defects identified.

The empty L0 collections follow the existing parametrization convention, L0 CI does not invoke this distributed test file, and the L1/L2 jobs continue to use valid nonempty shape sets.

Important Files Changed

Filename	Overview
tests/jax/test_distributed_fused_attn.py	Correctly converts invalid L0 shape entries into empty parameter sets without affecting the L1/L2 distributed cases, and adds focused regression coverage.

_{Reviews (1): Last reviewed commit: "fix: L0 placeholder shapes are invalid e..." | Re-trigger Greptile}

[cyanguwa]

I think the test is not really necessary - could you please remove it? Thanks.