fix: reorder test uses original shape after swapping SBHD axes

[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_distributed_fused_attn.py: reorder test uses original shape after swapping SBHD axes.

Changes

• tests/jax/test_distributed_fused_attn.py: reorder test uses original shape after swapping SBHD axes.

Details

--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -1,4 +1,4 @@
-        if reorder_strategy == ReorderStrategy.Striped:
-            seq_lens = shape[seq_dim]
-            if seq_lens < (cp_size * stripe_size):
-                pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")
+        if reorder_strategy == ReorderStrategy.Striped:
+            seq_lens = tensor.shape[seq_dim]
+            if seq_lens < (cp_size * stripe_size):
+                pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}")

Tests

• tests/jax/test_distributed_fused_attn.py

[greptile-apps]

Greptile Summary

This PR corrects the striped causal-load-balancing test guard to inspect the post-transformation tensor shape, preserving valid SBHD test coverage after its axes are swapped.

• Uses tensor.shape[seq_dim] when determining the effective sequence length.
• Prevents valid SBHD configurations from being skipped based on the original batch dimension.

Confidence Score: 5/5

The PR appears safe to merge because the updated guard consistently reads the effective sequence axis after SBHD axis swapping.

The changed lookup fixes erroneous SBHD skips, while the current BSHD, SBHD, and THD parameter combinations continue to satisfy the striped reorder preconditions.

Important Files Changed

Filename	Overview
tests/jax/test_distributed_fused_attn.py	Correctly derives striped reorder sequence length from the post-swap tensor; no changed-code defect was identified.

_{Reviews (1): Last reviewed commit: "fix: reorder test uses original shape af..." | Re-trigger Greptile}