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

Bump the shard count for TPU to avoid timeouts

Browse Source 
PiperOrigin-RevId: 487421018
This commit is contained in:
Yash Katariya 2022-11-09 20:31:32 -08:00 committed by jax authors
parent e42e52d4aa
commit 71360edf90
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tests/BUILD
View File

@ -384,7 +384,7 @@ jax_test(
shard_count = {
"cpu": 40,
"gpu": 40,
"tpu": 10,
"tpu": 40,
"iree": 10,
},
)
@ -452,7 +452,7 @@ jax_test(
shard_count = {
"cpu": 40,
"gpu": 40,
"tpu": 20,
"tpu": 40,
"iree": 40,
},
deps = [":lax_test_lib"],
@ -661,7 +661,7 @@ jax_test(
shard_count = {
"cpu": 40,
"gpu": 40,
"tpu": 40,
"tpu": 50,
},
)
@ -688,7 +688,7 @@ jax_test(
shard_count = {
"cpu": 40,
"gpu": 40,
"tpu": 40,
"tpu": 50,
"iree": 10,
},
deps = [