fix: make_mask uses deprecated jnp.bool instead of jnp.bool_

[andrewwhitecdw]

This PR addresses the following issue in tests/jax/test_fused_attn.py: make_mask uses deprecated jnp.bool instead of jnp.bool_.

Changes

• tests/jax/test_fused_attn.py: make_mask uses deprecated jnp.bool instead of jnp.bool_.

Details

--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1,13 +1,13 @@
-    # sliding window mask
-    inv_swa_mask = (
-        make_swa_mask(
-            segment_pos_q,
-            segment_pos_kv,
-            window_size,
-            dtype=jnp.bool,
-            segment_ids_q=segment_ids_q,
-            segment_ids_kv=segment_ids_kv,
-        )
-        if attn_mask_type.is_bottom_right()
-        else make_swa_mask(segment_pos_q, segment_pos_kv, window_size, dtype=jnp.bool_)
-    )
+    # sliding window mask
+    inv_swa_mask = (
+        make_swa_mask(
+            segment_pos_q,
+            segment_pos_kv,
+            window_size,
+            dtype=jnp.bool_,
+            segment_ids_q=segment_ids_q,
+            segment_ids_kv=segment_ids_kv,
+        )
+        if attn_mask_type.is_bottom_right()
+        else make_swa_mask(segment_pos_q, segment_pos_kv, window_size, dtype=jnp.bool_)
+    )

Tests

• tests/jax/test_fused_attn.py

--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -236,5 +236,22 @@
     return mask
     
     
+def test_make_mask_bottom_right_swa_dtype():
+    """Regression test: make_mask bottom-right branch should use jnp.bool_, not jnp.bool."""
+    batch, seqlen = 2, 16
+    segment_ids = jnp.ones((batch, seqlen), dtype=jnp.int32)
+    segment_pos = jnp.broadcast_to(jnp.arange(seqlen, dtype=jnp.int32), (batch, seqlen))
+    window_size = (4, 0)
+    mask = make_mask(
+        segment_ids,
+        segment_ids,
+        segment_pos,
+        segment_pos,
+        AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK,
+        window_size,
+    )
+    assert mask.dtype == jnp.bool_
+
+
 @jax.jit
 def get_seqlens_and_offsets(segment_ids):

[greptile-apps]

Greptile Summary

This PR replaces the deprecated jnp.bool alias with jnp.bool_ in the bottom-right sliding-window mask path.

• Updates the dtype passed to make_swa_mask.
• Adds focused coverage for the bottom-right mask path.

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	Replaces the deprecated JAX boolean alias and adds coverage for the affected bottom-right sliding-window mask branch.

_{Reviews (2): Last reviewed commit: "fix: make_mask uses deprecated jnp.bool ..." | Re-trigger Greptile}

[greptile-apps]

Assertion misses deprecated alias

This assertion checks only the final mask dtype, which is boolean after the logical mask operations whether make_swa_mask receives jnp.bool or jnp.bool_; restoring the deprecated argument therefore leaves this regression test passing.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[cyanguwa]

We seem to have a few other call sites of jnp.bool? If it's deprecated, could we fix those ones as well? Thanks.

I think we can get away with the test as well.