Import framework submodules lazily (#839)

* Import frameworks lazily
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Only load modules lazily after an import error

Pylint doesn't handle lazy loading gracefully.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

transformer_engine/__init__.py

@@ -3,15 +3,43 @@
 # See LICENSE for license information.
 
 """Top level package"""
+import importlib.util
+import sys
+from types import ModuleType
+
 from ._version import __version__
 from . import common
 
+def _lazy_import(name: str) -> ModuleType:
+    """Construct a module that is imported the first time it is used"""
+    spec = importlib.util.find_spec(name)
+    loader = importlib.util.LazyLoader(spec.loader)
+    spec.loader = loader
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    loader.exec_module(module)
+    return module
+
+# Import framework submodules
+# Note: Load module lazily if import fails. This way a useful import
+# error will be thrown if the user attempts to access the module.
 try:
     from . import pytorch
-except ImportError as e:
-    pass
-
+except ImportError:
+    pytorch = _lazy_import("transformer_engine.pytorch")
 try:
     from . import jax
-except ImportError as e:
-    pass
+except ImportError:
+    jax = _lazy_import("transformer_engine.jax")
+try:
+    from . import paddle
+except ImportError:
+    paddle = _lazy_import("transformer_engine.paddle")
+
+__all__ = [
+    "__version__",
+    "common",
+    "jax",
+    "paddle",
+    "pytorch",
+]