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

Increase sharding of or disable some flaky CI tests.

Browse Source 
PiperOrigin-RevId: 529405705
This commit is contained in:
Peter Hawkins 2023-05-04 07:41:21 -07:00 committed by jax authors
parent 40aa4e1781
commit 09fce87f54
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tests/BUILD
View File

@ -635,7 +635,7 @@ jax_test(
shard_count = {
"cpu": 30,
"gpu": 30,
"tpu": 30,
"tpu": 40,
"iree": 30,
},
tags = ["noasan"], # Times out
@ -652,7 +652,8 @@ jax_test(
"notsan",
],
"tpu": [
"noasan", # Times out under asan/tsan.
"noasan", # Times out under asan/msan/tsan.
"nomsan",
"notsan",
"optonly",
],
@ -719,6 +720,9 @@ jax_test(
jax_test(
name = "scipy_stats_test",
srcs = ["scipy_stats_test.py"],
backend_tags = {
"tpu": ["nomsan"], # Times out
},
shard_count = {
"cpu": 40,
"gpu": 30,
@ -1029,7 +1033,7 @@ jax_test(
shard_count = {
"cpu": 20,
"gpu": 10,
"tpu": 10,
"tpu": 20,
},
)
@ -1037,7 +1041,7 @@ jax_test(
name = "shard_map_test",
srcs = ["shard_map_test.py"],
shard_count = {
"cpu": 20,
"cpu": 30,
"gpu": 10,
"tpu": 10,
},