fix: make_mask uses deprecated jnp.bool instead of jnp.bool_ (#3371)

fix: use jnp.bool_ instead of deprecated jnp.bool in make_mask

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>

Author: andrewwhitecdw

tests/jax/test_fused_attn.py

@@ -222,7 +222,7 @@ def make_mask(
             segment_pos_q,
             segment_pos_kv,
             window_size,
-            dtype=jnp.bool,
+            dtype=jnp.bool_,
             segment_ids_q=segment_ids_q,
             segment_ids_kv=segment_ids_kv,
         )