From 85a466d730079a9daaa94486607816dc83701508 Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Tue, 24 Sep 2024 13:09:42 -0700 Subject: [PATCH] Lower the shard count for sparse_bcoo_bcsr_test on TPU as well. There are flaky timeouts in CI, and we've already lowered the shard count on multiple other platforms. PiperOrigin-RevId: 678367575 --- tests/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/BUILD b/tests/BUILD index 49dbf0512..0cc6ed6d9 100644 --- a/tests/BUILD +++ b/tests/BUILD @@ -997,6 +997,7 @@ jax_multiplatform_test( "cpu": ["--jax_num_generated_cases=40"], "cpu_x32": ["--jax_num_generated_cases=40"], "gpu": ["--jax_num_generated_cases=40"], + "tpu": ["--jax_num_generated_cases=40"], }, shard_count = { "cpu": 50,