Enable the bounds checker in the Pallas:SC bounds checker test

PiperOrigin-RevId: 802046061

Author: apaszke

tests/pallas/BUILD

@@ -1198,6 +1198,7 @@ jax_multiplatform_test(
     ] + py_deps([
         "numpy",
         "absl/testing",
+        "absl/flags",
     ]),
 )
 

tests/pallas/tpu_pallas_sparsecore_debug_check_test.py

@@ -28,6 +28,7 @@
 import unittest
 
 from absl.testing import absltest
+from absl import flags
 import jax
 from jax._src import test_util as jtu
 from jax.experimental import pallas as pl
@@ -118,6 +119,12 @@ def kernel(_):
     )
 
   def test_trigger_bounds_checker(self):
+    if "xla_sc_assert_level" in flags.FLAGS:
+      # The test crashes the process anyway, so no need to be clean.
+      flags.FLAGS.xla_sc_assert_level = "bounds"
+    else:
+      self.skipTest("TODO: Find another way to enable bounds checking.")
+
     x = jnp.arange(8, dtype=jnp.int32)
     # Index 8 is out-of-bounds.
     indices = jnp.array([0, 1, 2, 3, 4, 5, 6, 8], dtype=jnp.int32)