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

Update TPU test configuration tags.

Browse Source 
PiperOrigin-RevId: 672562923
This commit is contained in:
Peter Hawkins 2024-09-09 09:01:51 -07:00 committed by jax authors
parent c0dacbf724
commit 95f38d95d7
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/BUILD
View File

@ -235,8 +235,8 @@ jax_test(
},
enable_configs = [
"gpu_2gpu_shardy",
"tpu_df_2x2_shardy",
"tpu_pf_2x2_shardy",
"tpu_v3_2x2_shardy",
"tpu_v4_2x2_shardy",
],
shard_count = {
"cpu": 5,
@ -1426,7 +1426,7 @@ jax_test(
name = "export_test",
srcs = ["export_test.py"],
enable_configs = [
"tpu_df_2x2",
"tpu_v3_2x2",
],
tags = [],
deps = [