From 1d20d2f30121ac395b6db27f7463b5cf24ad4ef1 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Fri, 19 May 2023 07:44:18 -0700
Subject: [PATCH] Increase sharding of host_callback_test on TPU to fix CI
 flakiness.

PiperOrigin-RevId: 533451822
---
 tests/BUILD | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/BUILD b/tests/BUILD
index 88e437693..f11b8fec7 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -914,6 +914,9 @@ jax_test(
     name = "host_callback_test",
     srcs = ["host_callback_test.py"],
     args = ["--jax_host_callback_outfeed=true"],
+    shard_count = {
+        "tpu": 5,
+    },
     deps = [
         "//jax:experimental",
         "//jax:experimental_host_callback",