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

sparse_test: reduce num_generated_cases to avoid timeouts

Browse Source 
PiperOrigin-RevId: 509941080
This commit is contained in:
Jake VanderPlas 2023-02-15 14:57:54 -08:00 committed by jax authors
parent cd0533cab0
commit 6608242f95
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tests/BUILD
View File

@ -755,6 +755,13 @@ jax_test(
],
"tpu": ["optonly"],
},
# Use fewer cases to prevent timeouts.
backend_variant_args = {
"cpu": ["--jax_num_generated_cases=40"],
"cpu_x32": ["--jax_num_generated_cases=40"],
"cpu_no_jax_array": ["--jax_num_generated_cases=40"],
"gpu": ["--jax_num_generated_cases=40"],
},
shard_count = {
"cpu": 50,
"gpu": 50,