Skip to content

Commit

Permalink
Reduce the max_examples for splash attention tests
Browse files Browse the repository at this point in the history
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
  • Loading branch information
apaszke authored and Google-ML-Automation committed Dec 17, 2024
1 parent 3fc2371 commit e1f037b
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion tests/pallas/tpu_splash_attention_kernel_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down

0 comments on commit e1f037b

Please sign in to comment.