From f0bde75dd32d4cbcc0c7e396779b0cd51370041d Mon Sep 17 00:00:00 2001
From: Hyeontaek Lim <hyeontaek@google.com>
Date: Thu, 21 Sep 2023 22:52:00 -0700
Subject: [PATCH] [JAX] Export shard_map_test for testing on additional JAX
 backends

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

diff --git a/tests/BUILD b/tests/BUILD
index 68831170b..17204aba4 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-load("@rules_python//python:defs.bzl", "py_test")
 load(
     "//jaxlib:jax.bzl",
     "jax_generate_backend_suites",
@@ -21,6 +20,7 @@ load(
     "py_deps",
     "pytype_test",
 )
+load("@rules_python//python:defs.bzl", "py_test")
 
 licenses(["notice"])
 
@@ -1207,6 +1207,7 @@ exports_files(
         "pmap_test.py",
         "pjit_test.py",
         "python_callback_test.py",
+        "shard_map_test.py",
         "transfer_guard_test.py",
     ],
     visibility = jax_test_file_visibility,