[Pallas:SC] Fix the hypothesis OOM condition to avoid v6e OOMs

PiperOrigin-RevId: 835230881

Author: apaszke

tests/pallas/tpu_sparsecore_pallas_test.py

@@ -323,7 +323,8 @@ def test_block_spec_untiled_slicing(self, data):
       slice_shape[-1] *= 128
     else:
       slice_shape[-1] *= 8
-    hp.assume(math.prod(slice_shape) <= 25000)  # Avoid OOMs.
+    max_elems = 12000 if jtu.is_device_tpu(6, "e") else 25000
+    hp.assume(math.prod(slice_shape) <= max_elems)  # Avoid OOMs.
     rank = len(slice_shape)
     offsets = data.draw(
         hps.lists(hps.integers(0, 4), min_size=rank, max_size=rank)