#sdy Add CPU targets in JAX.

PiperOrigin-RevId: 678626718
jax-ml · Sep 26, 2024 · 9d9e468 · 9d9e468
1 parent e62a50c
commit 9d9e468
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/BUILD b/tests/BUILD
@@ -234,6 +234,7 @@ jax_multiplatform_test(
         "gpu": ["noasan"],  # Memory leaks in NCCL, see https://github.com/NVIDIA/nccl/pull/1143
     },
     enable_configs = [
+        "cpu_shardy",
         "gpu_2gpu_shardy",
         "tpu_v3_2x2_shardy",
         "tpu_v4_2x2_shardy",
@@ -1347,6 +1348,7 @@ jax_multiplatform_test(
     name = "shard_map_test",
     srcs = ["shard_map_test.py"],
     enable_configs = [
+        "cpu_shardy",
         "gpu_2gpu_shardy",
         "tpu_v3_2x2_shardy",
         "tpu_v4_2x2_shardy",