[Mosaic GPU] Disable XLA autotuning in matmul_test.py.

This slows down correctness tests for no value whatsoever. PiperOrigin-RevId: 637805645
jax-ml · May 28, 2024 · 72b111a · 72b111a
1 parent 97f9a5e
commit 72b111a
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/mosaic/matmul_test.py b/tests/mosaic/matmul_test.py
@@ -14,6 +14,8 @@
 # ==============================================================================
 """Test different parameterizations of a matmul."""
 
+import os
+
 from absl.testing import absltest, parameterized
 from jax._src import config
 from jax._src import test_util as jtu
@@ -28,6 +30,8 @@
 
 config.update("jax_traceback_filtering", "off")
 config.parse_flags_with_absl()
+os.environ["XLA_FLAGS"] = (
+    os.environ.get("XLA_FLAGS", "") + " --xla_gpu_autotune_level=0")
 
 
 class MatmulTestCase(jtu.JaxTestCase):