Increase some test shardings to reduce CI timeouts under asan/tsan/msan.

PiperOrigin-RevId: 469173324
jax-ml · Aug 22, 2022 · da4e79a · da4e79a
1 parent f905d98
commit da4e79a
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/BUILD b/tests/BUILD
@@ -317,9 +317,9 @@ jax_test(
  name = "lax_control_flow_test",
  srcs = ["lax_control_flow_test.py"],
  shard_count = {
- "cpu": 10,
+ "cpu": 20,
  "gpu": 20,
- "tpu": 10,
+ "tpu": 20,
  "iree": 10,
  },
 )
@@ -529,9 +529,9 @@ jax_test(
  name = "pmap_test",
  srcs = ["pmap_test.py"],
  shard_count = {
- "cpu": 15,
+ "cpu": 30,
  "gpu": 30,
- "tpu": 15,
+ "tpu": 30,
  },
  tags = ["multiaccelerator"],
  deps = [