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

Lower the shard count for sparse_bcoo_bcsr_test on TPU as well.

Browse Source 
There are flaky timeouts in CI, and we've already lowered the shard count on multiple other platforms.

PiperOrigin-RevId: 678367575
This commit is contained in:
Peter Hawkins 2024-09-24 13:09:42 -07:00 committed by jax authors
parent 70f91db853
commit 85a466d730
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tests/BUILD
View File

@ -997,6 +997,7 @@ jax_multiplatform_test(
"cpu": ["--jax_num_generated_cases=40"],
"cpu_x32": ["--jax_num_generated_cases=40"],
"gpu": ["--jax_num_generated_cases=40"],
"tpu": ["--jax_num_generated_cases=40"],
},
shard_count = {
"cpu": 50,