fix: Optional segment position annotations in mask helpers

[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_fused_attn.py: Optional segment position annotations in mask helpers.

Changes

• tests/jax/test_fused_attn.py: Optional segment position annotations in mask helpers.

Details

--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1,19 +1,19 @@
-@jax.jit
-def make_causal_mask(
-    segment_ids_q: ArrayLike,
-    segment_ids_kv: ArrayLike,
-    segment_pos_q: ArrayLike = None,
-    segment_pos_kv: ArrayLike = None,
-) -> Array:
-
-...
-
-@partial(jax.jit, static_argnums=(4, 5))
-def make_mask(
-    segment_ids_q: ArrayLike,
-    segment_ids_kv: ArrayLike,
-    segment_pos_q: ArrayLike,
-    segment_pos_kv: ArrayLike,
-    attn_mask_type: AttnMaskType,
-    window_size: Optional[Tuple[int, int]] = None,
-) -> Array:
+@jax.jit
+def make_causal_mask(
+    segment_ids_q: ArrayLike,
+    segment_ids_kv: ArrayLike,
+    segment_pos_q: Optional[ArrayLike] = None,
+    segment_pos_kv: Optional[ArrayLike] = None,
+) -> Array:
+
+...
+
+@partial(jax.jit, static_argnums=(4, 5))
+def make_mask(
+    segment_ids_q: ArrayLike,
+    segment_ids_kv: ArrayLike,
+    segment_pos_q: Optional[ArrayLike],
+    segment_pos_kv: Optional[ArrayLike],
+    attn_mask_type: AttnMaskType,
+    window_size: Optional[Tuple[int, int]] = None,
+) -> Array:

Tests

• tests/jax/test_fused_attn.py

--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -180,6 +180,28 @@
     mask = jnp.logical_not(inv_mask)
     return mask
 
+
+def test_mask_helpers_handle_none_segment_pos():
+    """Verify make_causal_mask and make_mask fall back to arange when segment_pos is None."""
+    seq = 8
+    segment_ids = jnp.array(
+        [[1, 1, 1, 1, 0, 0, 0, 0], [2, 2, 2, 2, 2, 0, 0, 0]], dtype=jnp.int32
+    )
+    segment_pos = jnp.broadcast_to(jnp.arange(seq, dtype=jnp.int32), segment_ids.shape)
+
+    # make_causal_mask uses default None for segment positions.
+    mask_with_pos = make_causal_mask(segment_ids, segment_ids, segment_pos, segment_pos)
+    mask_without_pos = make_causal_mask(segment_ids, segment_ids)
+    assert mask_with_pos.shape == mask_without_pos.shape
+    assert jnp.array_equal(mask_with_pos, mask_without_pos)
+
+    # make_mask accepts explicit None for segment positions.
+    for mask_type in (AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK):
+        mask_with_pos = make_mask(
+            segment_ids, segment_ids, segment_pos, segment_pos, mask_type
+        )
+        mask_without_pos = make_mask(segment_ids, segment_ids, None, None, mask_type)
+        assert mask_with_pos.shape == mask_without_pos.shape
+        assert jnp.array_equal(mask_with_pos, mask_without_pos)
+
 
 @jax.jit
 def get_seqlens_and_offsets(segment_ids):

[greptile-apps]

Greptile Summary

Corrects the BRCM non-zero-dropout skip message in the JAX fused-attention test runner.

• Fixes the spelling of “dropouts.”
• Removes an unnecessary f-string prefix from the static message.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
tests/jax/test_fused_attn.py	Corrects a static pytest skip message without changing test behavior or runtime logic.

_{Reviews (2): Last reviewed commit: "fix: Optional segment position annotatio..." | Re-trigger Greptile}

[cyanguwa]

The code change doesn't match the PR title, and is duplicate to PR 3370.