Repository: jax-ml/jax
Pull request #25424 (2024-12-11 22:30:46)

internal: dedupe lax broadcasting logic


Patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 221d4a4fb..7ed7fb46e 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -102,27 +102,31 @@ def _validate_shapes(shapes: Sequence[Shape]):
   else:
     map(_check_static_shape, shapes)
 
-def _try_broadcast_shapes(
-    shapes: Sequence[tuple[int, ...]]) -> tuple[int, ...] | None:
-  if len(shapes) == 1: return shapes[0]
+def _try_broadcast_shapes(*shapes: tuple[int, ...], name: str) -> tuple[int, ...]:
+  """
+  Attempt to broadcast shapes, raising a TypeError if broadcasting fails.
+  """
+  if not shapes:
+    raise TypeError(f"{name}: At least one shape is required.")
   ranks = {len(shape) for shape in shapes}
-  if len(ranks) > 1: return None  # must have consistent rank
-  rank = ranks.pop()
-  if not rank: return ()  # scalar case
+  if len(ranks) != 1:
+    raise TypeError(f'{name}: arrays must have the same number of dimensions,'
+                    f' got {ranks}')
   result_shape = []
-  for ds in unsafe_zip(*shapes):
+  for ds in zip(*shapes):
     if all(core.same_referent(d, ds[0]) for d in ds[1:]):
       # if all axes are identical objects, the resulting size is the object
       result_shape.append(ds[0])
     else:
-      # if all dims are equal (or 1), the result is the non-1 size (or 1)
+      # if all dims are equal (or 1), the result is the non-1 size
       non_1s = [d for d in ds if not core.definitely_equal(d, 1)]
       if not non_1s:
         result_shape.append(1)
       elif all(core.definitely_equal(non_1s[0], d) for d in non_1s[1:]):
         result_shape.append(non_1s[0])
       else:
-        return None
+        raise TypeError(f'{name} got incompatible shapes for broadcasting: '
+                        f'{", ".join(map(str, map(tuple, shapes)))}.')
   return tuple(result_shape)
 
 def asarray(x: ArrayLike) -> Array:
@@ -165,11 +169,12 @@ def _broadcast_shapes_uncached(*shapes):
     return result_shape
 
   # Next try singleton-broadcasting, padding out ranks using singletons.
-  shape_list = [(1,) * (ndim - len(shape)) + shape for shape in shapes]
-  result_shape = _try_broadcast_shapes(shape_list)
-  if result_shape is None:
-    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}")
-  return result_shape
+  rank_promoted_shapes = tuple((*((1,) * (ndim - len(shape))), *shape) for shape in shapes)
+  try:
+    return _try_broadcast_shapes(*rank_promoted_shapes, name='broadcast_shapes')
+  except TypeError as err:
+    # Raise ValueError here for backward compatibility.
+    raise ValueError(f"Incompatible shapes for broadcasting: shapes={list(shapes)}") from err
 
 def _identity(x): return x
 
@@ -2133,27 +2138,7 @@ def broadcasting_shape_rule(name, *avals):
   shapes = [aval.shape for aval in avals if aval.shape]
   if not shapes:
     return ()
-  if len({len(shape) for shape in shapes}) != 1:
-    msg = '{}: arrays must have same number of dimensions, got {}.'
-    raise TypeError(msg.format(name, ', '.join(map(str, map(tuple, shapes)))))
-  # TODO(mattjj): de-duplicate with _try_broadcast_shapes
-  result_shape = []
-  for ds in zip(*shapes):
-    if all(core.same_referent(d, ds[0]) for d in ds[1:]):
-      # if all axes are identical objects, the resulting size is the object
-      result_shape.append(ds[0])
-    else:
-      # if all dims are equal (or 1), the result is the non-1 size
-      non_1s = [d for d in ds if not core.definitely_equal(d, 1)]
-      if not non_1s:
-        result_shape.append(1)
-      elif all(core.definitely_equal(non_1s[0], d) for d in non_1s[1:]):
-        result_shape.append(non_1s[0])
-      else:
-        raise TypeError(f'{name} got incompatible shapes for broadcasting: '
-                        f'{", ".join(map(str, map(tuple, shapes)))}.')
-
-  return tuple(result_shape)
+  return _try_broadcast_shapes(*shapes, name=name)
 
 
 def broadcasting_sharding_rule(name, *avals):
diff --git a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py
index 4f08e5865..7124d64f0 100644
--- a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py
+++ b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py
@@ -603,9 +603,9 @@ def _apply_mask_and_soft_cap(
     else:
       mask = pl.load(mask_ref, (k_slice, slice(None)))
 
-    masks.append(
-        jnp.bitwise_or(mask, jnp.broadcast_to(should_not_mask, mask.shape))
-    )
+    snm = jnp.where(should_not_mask, 1, 0)
+    masks.append(jnp.bitwise_or(mask, jnp.broadcast_to(snm, mask.shape)) != 0)
+
   if mask_function is not None:
     # Compute the mask using the given q_sequence indices.
     # KV indices are computed on the fly. This works because we only support Q
@@ -900,16 +900,6 @@ def _splash_attention_forward(
     kv_seq_len_dimension = 1
     num_kv_heads = k.shape[0]
 
-  partial_mask_blocks = fwd_mask_info.partial_mask_blocks
-  if (
-      partial_mask_blocks is not None
-      and jnp.dtype(partial_mask_blocks.dtype) != np.bool_
-  ):
-    raise ValueError(
-        "partial_mask_blocks must be of type np.bool_ but got"
-        f" {partial_mask_blocks.dtype}"
-    )
-
   if len(k.shape) != expected_kv_rank:
     raise ValueError(
         f"Expected {expected_kv_rank}-dim 'key' tensor for MQA. Instead got a"
diff --git a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py
index a5a524f70..99ab04898 100644
--- a/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py
+++ b/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py
@@ -56,7 +56,7 @@ class MaskInfo(NamedTuple):
       indicates that the corresponding block in the full mask contained both
       zeros and ones. An entry of 2 indicates the corresponding block was
       entirely ones.
-    partial_mask_blocks: A bool[num_partial_blocks, block_q, block_kv] NumPy
+    partial_mask_blocks: A i32[num_partial_blocks, block_q, block_kv] NumPy
       array that contains the blocks of the original mask that contained both
       zeros and ones. The entries in `mask_next` point to indices in the first
       axis of this array.
@@ -305,7 +305,7 @@ def _get_mask_info_for_shard(
 
 
 # When used in a transformer network with multiple layers, the SplashAttention
-# kernel is created several times with the same mask. Cache MaskInfo to avoid
+# kernel is created serveral times with the same mask. Cache MaskInfo to avoid
 # blowing up compile times. Ideally the size of the cache should be determined
 # by the client.
 @functools.lru_cache(maxsize=12)
@@ -376,6 +376,14 @@ def _process_mask(
   if mod != 0:
     raise ValueError(f'{head_shards=} should divide {head_count=}.')
 
+  first_mask_size = mask.masks[0].shape
+  for h in range(head_count):
+    if mask.masks[h].shape != first_mask_size:
+      raise ValueError(
+          f'First head mask has shape {first_mask_size}, but head mask {h} has'
+          f' shape {mask.masks[h].shape}. All head masks must have the same'
+          ' shape.'
+      )
 
   # Uniquify the masks.
   # Create a collection of the unique head masks in the input multi-head mask.
@@ -518,9 +526,13 @@ def _process_mask(
 
   partial_mask_blocks = None
   has_mask_next = False
-  if len(unique_partial_mask_blocks) >= 1:
+  if len(unique_partial_mask_blocks) == 1:
+    partial_mask_blocks = [x.array for x in unique_partial_mask_blocks]
+    partial_mask_blocks = partial_mask_blocks[0][None].astype(np.int32)
+    has_mask_next = True
+  elif len(unique_partial_mask_blocks) > 1:
     partial_mask_blocks = [x.array for x in unique_partial_mask_blocks]
-    partial_mask_blocks = np.stack(partial_mask_blocks, axis=0).astype(np.bool_)
+    partial_mask_blocks = np.stack(partial_mask_blocks, axis=0).astype(np.int32)
     has_mask_next = True
   if is_dkv and partial_mask_blocks is not None:
     partial_mask_blocks = np.swapaxes(partial_mask_blocks, -1, -2)


Test patch:
diff --git a/tests/pallas/tpu_splash_attention_mask_test.py b/tests/pallas/tpu_splash_attention_mask_test.py
index ab4f68757..ce7d8fd09 100644
--- a/tests/pallas/tpu_splash_attention_mask_test.py
+++ b/tests/pallas/tpu_splash_attention_mask_test.py
@@ -18,9 +18,9 @@ from __future__ import annotations
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
-from jax._src import test_util as jtu
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask as mask_lib
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask_info as mask_info_lib
+from jax._src import test_util as jtu
 import numpy as np
 
 jax.config.parse_flags_with_absl()
@@ -798,7 +798,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
         self._expected_causal_data_next[None],
         self._expected_causal_mask_next(0)[None] if not is_lazy_mask else None,
         self._expected_causal_block_mask[None],
-        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.bool_)), 0)
+        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.int32)), 0)
         if not is_lazy_mask
         else None,
         np.arange(sequence_lengths[0], dtype=np.int32)
@@ -813,7 +813,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
         else None,
         self._expected_causal_block_mask_dkv[None],
         np.expand_dims(
-            np.tril(np.ones(block_shape, dtype=np.bool_)), 0
+            np.tril(np.ones(block_shape, dtype=np.int32)), 0
         ).swapaxes(-1, -2)
         if not is_lazy_mask
         else None,
@@ -851,7 +851,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
         self._expected_causal_data_next[None],
         self._expected_causal_mask_next(0)[None] if not is_lazy_mask else None,
         self._expected_causal_block_mask[None],
-        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.bool_)), 0)
+        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.int32)), 0)
         if not is_lazy_mask
         else None,
         np.arange(sequence_lengths[0], dtype=np.int32)
@@ -894,7 +894,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
         expected_causal_mask_next_dkv if not is_lazy_mask else None,
         expected_causal_block_mask_dkv,
         np.expand_dims(
-            np.tril(np.ones(block_shape, dtype=np.bool_)), 0
+            np.tril(np.ones(block_shape, dtype=np.int32)), 0
         ).swapaxes(-1, -2)
         if not is_lazy_mask
         else None,
@@ -974,7 +974,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
         expected_causal_data_next,
         expected_causal_mask_next if not is_lazy_mask else None,
         expected_causal_block_mask,
-        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.bool_)), 0)
+        np.expand_dims(np.tril(np.ones(block_shape, dtype=np.int32)), 0)
         if not is_lazy_mask
         else None,
         np.arange(sequence_lengths[0], dtype=np.int32)
@@ -1029,7 +1029,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
         expected_causal_mask_next_dkv if not is_lazy_mask else None,
         expected_causal_block_mask_dkv,
         np.expand_dims(
-            np.tril(np.ones(block_shape, dtype=np.bool_)), 0
+            np.tril(np.ones(block_shape, dtype=np.int32)), 0
         ).swapaxes(-1, -2)
         if not is_lazy_mask
         else None,
@@ -1069,10 +1069,10 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     expected_partial_mask_blocks = self._stack(
         [
             np.triu(
-                np.tri(*block_shape, window_size, dtype=np.bool_), -window_size
+                np.tri(*block_shape, window_size, dtype=np.int32), -window_size
             ),
-            np.tri(*block_shape, -window_size, dtype=np.bool_),
-            np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+            np.tri(*block_shape, -window_size, dtype=np.int32),
+            np.triu(np.ones(block_shape, dtype=np.int32), window_size),
         ],
     )
 
@@ -1179,8 +1179,8 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
 
     expected_partial_mask_blocks = self._stack(
         [
-            np.triu(np.tri(*block_shape, 0, dtype=np.bool_), -window_size),
-            np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+            np.triu(np.tri(*block_shape, 0, dtype=np.int32), -window_size),
+            np.triu(np.ones(block_shape, dtype=np.int32), window_size),
         ],
     )
 
@@ -1298,13 +1298,13 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     )
 
     expected_partial_mask_blocks = self._stack([
-        np.tril(np.ones(block_shape, dtype=np.bool_)),
+        np.tril(np.ones(block_shape, dtype=np.int32)),
         np.triu(
-            np.tri(*block_shape, window_size, dtype=np.bool_),
+            np.tri(*block_shape, window_size, dtype=np.int32),
             -window_size,
         ),
-        np.tri(*block_shape, -window_size, dtype=np.bool_),
-        np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
     ])
 
     expected_block_mask_dkv = self._stack(
@@ -1384,7 +1384,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     ])
 
     expected_partial_mask_blocks = np.expand_dims(
-        np.tril(np.ones(block_shape, dtype=np.bool_)), 0
+        np.tril(np.ones(block_shape, dtype=np.int32)), 0
     )
 
     expected_mask_info = mask_info_lib.MaskInfo(
@@ -1460,13 +1460,13 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     )
 
     expected_partial_mask_blocks = self._stack([
-        np.tril(np.ones(block_shape, dtype=np.bool_)),
+        np.tril(np.ones(block_shape, dtype=np.int32)),
         np.triu(
-            np.tri(*block_shape, window_size, dtype=np.bool_),
+            np.tri(*block_shape, window_size, dtype=np.int32),
             -window_size,
         ),
-        np.tri(*block_shape, -window_size, dtype=np.bool_),
-        np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
     ])
 
     expected_mask_info = mask_info_lib.MaskInfo(
@@ -1577,13 +1577,13 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     )
 
     expected_partial_mask_blocks = self._stack([
-        np.tril(np.ones(block_shape, dtype=np.bool_)),
+        np.tril(np.ones(block_shape, dtype=np.int32)),
         np.triu(
-            np.tri(*block_shape, window_size, dtype=np.bool_),
+            np.tri(*block_shape, window_size, dtype=np.int32),
             -window_size,
         ),
-        np.tri(*block_shape, -window_size, dtype=np.bool_),
-        np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
     ])
 
     expected_mask_info = mask_info_lib.MaskInfo(
@@ -1749,13 +1749,13 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     expected_partial_mask_blocks = self._stack([
         # Wide
         np.triu(
-            np.tri(*block_shape, window_size, dtype=np.bool_),
+            np.tri(*block_shape, window_size, dtype=np.int32),
             -window_size,
         ),
-        np.tri(*block_shape, -window_size, dtype=np.bool_),
-        np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+        np.tri(*block_shape, -window_size, dtype=np.int32),
+        np.triu(np.ones(block_shape, dtype=np.int32), window_size),
         # Narrow
-        np.triu(np.tri(*block_shape, 0, dtype=np.bool_), -window_size),
+        np.triu(np.tri(*block_shape, 0, dtype=np.int32), -window_size),
     ])
 
     expected_mask_info = mask_info_lib.MaskInfo(
@@ -1890,7 +1890,7 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
     )
 
     expected_partial_mask_blocks = np.expand_dims(
-        np.tril(np.ones(block_shape, dtype=np.bool_)), 0
+        np.tril(np.ones(block_shape, dtype=np.int32)), 0
     )
 
     expected_mask_info = mask_info_lib.MaskInfo(
@@ -1979,13 +1979,13 @@ class SplashAttentionMaskInfoTest(jtu.JaxTestCase):
 
     expected_partial_mask_blocks = self._stack(
         [
-            np.tril(np.ones(block_shape, dtype=np.bool_)),
+            np.tril(np.ones(block_shape, dtype=np.int32)),
             np.triu(
-                np.tri(*block_shape, window_size, dtype=np.bool_),
+                np.tri(*block_shape, window_size, dtype=np.int32),
                 -window_size,
             ),
-            np.tri(*block_shape, -window_size, dtype=np.bool_),
-            np.triu(np.ones(block_shape, dtype=np.bool_), window_size),
+            np.tri(*block_shape, -window_size, dtype=np.int32),
+            np.triu(np.ones(block_shape, dtype=np.int32), window_size),
         ],
     )