mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reduce the max_examples for splash attention tests

Browse Source 
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
This commit is contained in:
Adam Paszke 2024-12-17 09:05:27 -08:00 committed by jax authors
parent 3fc237125b
commit e1f037bffd
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/pallas/tpu_splash_attention_kernel_test.py
View File

@ -40,7 +40,7 @@ except (ModuleNotFoundError, ImportError):
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])