From 265b39d23f2e89b83ac52a97f01b2b95554c2c1f Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Mon, 26 Sep 2022 14:17:43 -0700
Subject: [PATCH] Add pytype_srcs to main jax BUILD rule

PiperOrigin-RevId: 476989241
---
 jax/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/jax/BUILD b/jax/BUILD
index 70df4fb3de7a..d473fbb62c10 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -110,6 +110,7 @@ py_library_providing_imports_info(
         "experimental/compilation_cache/cache_interface.py",
     ],
     lib_rule = pytype_library,
+    pytype_srcs = glob(["_src/**/*.pyi"]),
     visibility = ["//visibility:public"],
     deps = select({
                ":enable_jaxlib_build": [":jaxlib_deps"],