Enable Lazy loading

haiku/BUILD

@@ -34,6 +34,7 @@ py_library(
         "pad.py",
         "testing.py",
     ],
+    lazy_imports = True,
     strict_deps = False,
     visibility = ["//visibility:public"],
     deps = [

haiku/_src/BUILD

@@ -6,6 +6,7 @@ package(
         "//learning/deepmind/jax/haiku/metadata:__subpackages__",
         "//haiku:__subpackages__",
     ],
+    features = ["-py_lazy_imports"],
 )
 
 licenses(["notice"])

haiku/_src/integration/BUILD

@@ -3,6 +3,7 @@ load("//haiku/_src:build_defs.bzl", "hk_py_binary", "hk_py_library", "hk_py_test
 package(
     default_applicable_licenses = ["//haiku:license"],
     default_visibility = ["//visibility:private"],
+    features = ["-py_lazy_imports"],
 )
 
 licenses(["notice"])
@@ -260,6 +261,7 @@ hk_py_test(
         # pip: absl/testing:absltest
         # pip: absl/testing:parameterized
         "//haiku",
+        "//haiku/_src:test_utils",
     ],
 )
 

haiku/_src/integration/descriptors_test.py

@@ -19,6 +19,10 @@
 from haiku._src import test_utils
 from haiku._src.integration import descriptors
 
+# Explicitly import all modules to ensure they are loaded for the test.
+for _ in test_utils.find_internal_python_modules(hk):
+  pass
+
 ALL_MODULES = descriptors.ALL_MODULES
 IGNORED_MODULES = descriptors.IGNORED_MODULES
 

haiku/_src/integration/doctest_test.py

@@ -70,7 +70,7 @@ def run_test():
     # Many tests expect to be run as part of an `hk.transform`. We loop over all
     # exported symbols and run them in their own `hk.transform` so parameter and
     # module names don't clash.
-    for name in module.__all__:
+    for name in dir(module):
       test_names = []
 
       value = getattr(module, name)

haiku/_src/integration/shim_hk_test.py

@@ -20,6 +20,11 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import haiku as hk
+from haiku._src import test_utils
+
+# Explicitly import all modules to ensure they are loaded for the test.
+for _ in test_utils.find_internal_python_modules(hk):
+  pass
 
 
 def named_internal_modules():

haiku/_src/no_flax_test.py

@@ -44,7 +44,7 @@ def test_import_shim(self):
 
     # Errors are only thrown when users try to use flax.
     with self.assertRaisesRegex(ImportError, "lift.*require `flax`"):
-      hk.experimental.flax.lift  # pylint: disable=pointless-statement
+      print(hk.experimental.flax.lift)
 
   def test_inspect_stack_does_not_fail(self):
     self.assertNotEmpty(inspect.stack())