internal changes

jax/BUILD

@@ -153,6 +153,7 @@ py_library_providing_imports_info(
         ],
         # TODO(dsuo): Consider moving these files out of experimental if they're in the public API.
     ) + ["//jax/experimental:jax_public"],
+    lazy_imports = True,
     lib_rule = pytype_library,
     pytype_srcs = glob(
         [

jaxlib/jax.bzl

@@ -136,13 +136,19 @@ jax_extra_deps = []
 jax_gpu_support_deps = []
 jax2tf_deps = []
 
+def add_lazy_imports_flag(kwargs):
+    new_kwargs = dict(kwargs.items() + {"lazy_imports": True}.items())
+    return new_kwargs
+
 def pytype_library(name, pytype_srcs = None, **kwargs):
     _ = pytype_srcs  # @unused
-    py_library(name = name, **kwargs)
+    new_kwargs = add_lazy_imports_flag(kwargs)
+    py_library(name = name, **new_kwargs)
 
 def pytype_strict_library(name, pytype_srcs = [], **kwargs):
     data = pytype_srcs + (kwargs["data"] if "data" in kwargs else [])
     new_kwargs = {k: v for k, v in kwargs.items() if k != "data"}
+    new_kwargs = add_lazy_imports_flag(new_kwargs)
     py_library(name = name, data = data, **new_kwargs)
 
 py_strict_library = py_library
@@ -151,6 +157,7 @@ py_strict_test = py_test
 def py_library_providing_imports_info(*, name, lib_rule = py_library, pytype_srcs = [], **kwargs):
     data = pytype_srcs + (kwargs["data"] if "data" in kwargs else [])
     new_kwargs = {k: v for k, v in kwargs.items() if k != "data"}
+    new_kwargs = add_lazy_imports_flag(new_kwargs)
     lib_rule(name = name, data = data, **new_kwargs)
 
 def py_extension(name, srcs, copts, deps, linkopts = []):