fix: SBHD reorder skip uses original shape instead of swapped tensor

tests/jax/test_distributed_fused_attn.py

@@ -668,7 +668,7 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size):
             seq_dim = 0
 
         if reorder_strategy == ReorderStrategy.Striped:
-            seq_lens = shape[seq_dim]
+            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=}")
 
@@ -681,3 +681,20 @@ def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size):
         inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim, stripe_size)
 
         assert jnp.array_equal(inversed, ref)
+
+    @pytest.mark.parametrize("stripe_size", [1, 4])
+    def test_sbhd_striped_uses_swapped_seq_dim(self, stripe_size):
+        """Regression test: SBHD Striped skip must use the swapped sequence dim."""
+        cp_size = 2
+        shape = (1, 16, 1, 1)  # original [batch, seq, heads, dim]
+        tensor = random.normal(random.PRNGKey(42), shape, dtype=jnp.bfloat16)
+        tensor = tensor.swapaxes(0, 1)  # SBHD: [seq, batch, heads, dim]
+
+        # Old logic read original shape[0]=1 (batch) and skipped; seq_len after swap is 16.
+        reorder = jax.jit(reorder_causal_load_balancing, static_argnums=[1, 2, 3, 4])
+        inverse = jax.jit(inverse_reorder_causal_load_balancing, static_argnums=[1, 2, 3, 4])
+
+        reordered = reorder(tensor, ReorderStrategy.Striped, cp_size, 0, stripe_size)
+        inversed = inverse(reordered, ReorderStrategy.Striped, cp_size, 0, stripe_size)
+
+        assert jnp.array_equal(inversed, tensor)