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 in the bottom-right sliding-window mask path and adds a focused test.

• Uses jnp.bool_ consistently when constructing sliding-window masks.
• Adds coverage for the bottom-right padding-causal mask path, although its final-dtype assertion does not distinguish the fixed and deprecated arguments.

Confidence Score: 4/5

The PR appears safe to merge, with the non-blocking caveat that its new regression test does not actually detect restoration of the deprecated alias.

The production test helper now uses the supported boolean dtype alias, while the only accepted concern is ineffective regression coverage because the asserted final mask dtype is boolean in either implementation.

Files Needing Attention: tests/jax/test_fused_attn.py

Important Files Changed

Filename	Overview
tests/jax/test_fused_attn.py	The alias replacement is correct, but the added assertion would still pass if the deprecated alias were restored.

_{Reviews (1): 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.