From e1f037bffd6420d64af47aa853e7d2347c9e8004 Mon Sep 17 00:00:00 2001 From: Adam Paszke Date: Tue, 17 Dec 2024 09:05:27 -0800 Subject: [PATCH] Reduce the max_examples for splash attention tests They are too expensive to run all the time. We have 25 test methods, so it's ok to drop the number of examples each one of them covers. PiperOrigin-RevId: 707124466 --- tests/pallas/tpu_splash_attention_kernel_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pallas/tpu_splash_attention_kernel_test.py b/tests/pallas/tpu_splash_attention_kernel_test.py index e128d75c8ff2..317c84a56ed6 100644 --- a/tests/pallas/tpu_splash_attention_kernel_test.py +++ b/tests/pallas/tpu_splash_attention_kernel_test.py @@ -40,7 +40,7 @@ raise unittest.SkipTest("these tests require hypothesis") jax.config.parse_flags_with_absl() -jtu.setup_hypothesis() +jtu.setup_hypothesis(max_examples=5) partial = functools.partial Draw = TypeVar("Draw", bound=Callable[[hps.SearchStrategy[Any]], Any])