Lazily load nvmath subpackages and bindings wrappers

nvmath/__init__.py

@@ -2,31 +2,29 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-import importlib.metadata
-
+"""NVIDIA Math Python libraries."""
 
-# Attempt to preload libraries.  Fail silently if preload fails.
-def _force_lib_load():
-    from nvmath._utils import module_init_force_cupy_lib_load
-
-    module_init_force_cupy_lib_load()
-
-
-_force_lib_load()
+import importlib
+import importlib.metadata
 
-from nvmath import (  # noqa: E402
-    bindings,  # noqa: E402
-    fft,
-    linalg,
-    sparse,
-    tensor,
-)
-from nvmath._utils import (  # noqa: E402
-    ComputeType,
-    CudaDataType,
-    LibraryPropertyType,
-)
-from nvmath.memory import BaseCUDAMemoryManager, BaseCUDAMemoryManagerAsync, MemoryPointer  # noqa: E402
+# High-level subpackages (fft, linalg, sparse, tensor, bindings) and public
+# names from nvmath._utils / nvmath.memory are loaded on first access. Eager
+# imports here would pull in every library wrapper whenever a nested module
+# such as nvmath.bindings.cycurand is imported (the CuPy bindings-build path).
+# Native-library preload (module_init_force_cupy_lib_load) is likewise deferred
+# until a high-level subpackage is used; it must not run for cycurand-only
+# imports.
+_LAZY_MODULES = frozenset({"bindings", "fft", "linalg", "sparse", "tensor"})
+_LAZY_ATTRS = {
+    "BaseCUDAMemoryManager": "nvmath.memory",
+    "BaseCUDAMemoryManagerAsync": "nvmath.memory",
+    "ComputeType": "nvmath._utils",
+    "CudaDataType": "nvmath._utils",
+    "LibraryPropertyType": "nvmath._utils",
+    "MemoryPointer": "nvmath.memory",
+}
+_PRELOAD_ON = frozenset({"fft", "linalg", "sparse", "tensor"})
+_libs_preloaded = False
 
 __all__ = [
     "BaseCUDAMemoryManager",
@@ -42,4 +40,34 @@ def _force_lib_load():
     "tensor",
 ]
 
+
+def _ensure_libs_preloaded() -> None:
+    global _libs_preloaded
+    if _libs_preloaded:
+        return
+    _libs_preloaded = True
+    from nvmath._utils import module_init_force_cupy_lib_load
+
+    module_init_force_cupy_lib_load()
+
+
+def __getattr__(name: str):
+    if name in _LAZY_MODULES:
+        if name in _PRELOAD_ON:
+            _ensure_libs_preloaded()
+        mod = importlib.import_module(f"{__name__}.{name}")
+        globals()[name] = mod
+        return mod
+    origin = _LAZY_ATTRS.get(name)
+    if origin is not None:
+        value = getattr(importlib.import_module(origin), name)
+        globals()[name] = value
+        return value
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | _LAZY_MODULES | _LAZY_ATTRS.keys())
+
+
 __version__ = importlib.metadata.version("nvmath-python")

nvmath/bindings/__init__.py

@@ -4,41 +4,49 @@
 
 # type: ignore
 
-from nvmath.bindings import cublas
-from nvmath.bindings import cublasLt
-from nvmath.bindings import cudss
-from nvmath.bindings import cufft
-from nvmath.bindings import curand
-from nvmath.bindings import cusolver
-from nvmath.bindings import cusolverDn
-from nvmath.bindings import cusolverSp
-from nvmath.bindings import cusparse
-from nvmath.bindings import cusparseLt
-from nvmath.bindings import cutensor
-
-try:
-    # cufftMp is Linux-only.
-    from nvmath.bindings import cufftMp
-except ImportError:
-    cufftMp = None
-
-try:
-    # nvshmem is Linux-only.
-    from nvmath.bindings import nvshmem
-except ImportError:
-    nvshmem = None
-
-try:
-    # cublasMp is Linux-only.
-    from nvmath.bindings import cublasMp
-except ImportError:
-    cublasMp = None
-
-try:
-    # cusolverMp is Linux-only.
-    from nvmath.bindings import cusolverMp
-except ImportError:
-    cusolverMp = None
+import importlib
+
+# Load library wrappers on first access (PEP 562). Eager imports here would
+# map every wrapper .so (including optional *Mp / NVSHMEM modules) as soon as
+# any nvmath.bindings Cython submodule is loaded — for example CuPy's
+# `from nvmath.bindings.cycurand cimport ...`.
+_REQUIRED = frozenset(
+    {
+        "cublas",
+        "cublasLt",
+        "cudss",
+        "cufft",
+        "curand",
+        "cusolver",
+        "cusolverDn",
+        "cusolverSp",
+        "cusparse",
+        "cusparseLt",
+        "cutensor",
+        "nvpl",
+    }
+)
+_OPTIONAL = frozenset({"cufftMp", "nvshmem", "cublasMp", "cusolverMp"})
+_SUBMODULES = _REQUIRED | _OPTIONAL
+
+
+def __getattr__(name):
+    if name in _SUBMODULES:
+        try:
+            mod = importlib.import_module(f"{__name__}.{name}")
+        except ImportError:
+            if name in _OPTIONAL:
+                globals()[name] = None
+                return None
+            raise
+        globals()[name] = mod
+        return mod
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(set(globals()) | _SUBMODULES)
+
 
 __all__ = [
     "cublas",

nvmath/bindings/_internal/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# type: ignore
+
+import importlib
+
+
+def __getattr__(name):
+    # Submodules were previously imported as a side effect of eagerly loading
+    # every nvmath.bindings wrapper. Keep getattr() working now that wrappers
+    # are loaded on demand (see nvmath._utils.module_init_force_cupy_lib_load).
+    try:
+        mod = importlib.import_module(f"{__name__}.{name}")
+    except ImportError as exc:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc
+    globals()[name] = mod
+    return mod

nvmath/fft/__init__.py

@@ -2,6 +2,10 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from ._configuration import *  # noqa: F403
-from ._helpers import *  # noqa: F403
-from .fft import *  # noqa: F403
+from nvmath._utils import module_init_force_cupy_lib_load
+
+module_init_force_cupy_lib_load()
+
+from ._configuration import *  # noqa: E402, F403
+from ._helpers import *  # noqa: E402, F403
+from .fft import *  # noqa: E402, F403

nvmath/linalg/__init__.py

@@ -2,10 +2,13 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from nvmath._utils import module_init_force_cupy_lib_load
 from nvmath.bindings.cublas import ComputeType  # type: ignore
 
-from . import advanced
-from .generic import (
+module_init_force_cupy_lib_load()
+
+from . import advanced  # noqa: E402
+from .generic import (  # noqa: E402
     DiagonalMatrixQualifier,
     DiagType,
     ExecutionCPU,
@@ -23,7 +26,7 @@
     matmul,
     matrix_qualifiers_dtype,
 )
-from .generic.solvermod import (
+from .generic.solvermod import (  # noqa: E402
     DirectSolver,
     DirectSolverOptions,
     InvalidDirectSolverState,

nvmath/sparse/__init__.py

@@ -2,24 +2,28 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+from nvmath._utils import module_init_force_cupy_lib_load
+
+module_init_force_cupy_lib_load()
+
 from . import advanced, ust  # noqa: E402
-from .generic import ComputeType, ExecutionCUDA, Matmul, MatmulOptions, matmul
-from .generic import (
+from .generic import ComputeType, ExecutionCUDA, Matmul, MatmulOptions, matmul  # noqa: E402
+from .generic import (  # noqa: E402
     compile_add as compile_matmul_add,
 )
-from .generic import (
+from .generic import (  # noqa: E402
     compile_atomic_add as compile_matmul_atomic_add,
 )
-from .generic import (
+from .generic import (  # noqa: E402
     compile_epilog as compile_matmul_epilog,
 )
-from .generic import (
+from .generic import (  # noqa: E402
     compile_mul as compile_matmul_mul,
 )
-from .generic import (
+from .generic import (  # noqa: E402
     compile_prolog as compile_matmul_prolog,
 )
-from .generic import matrix_qualifiers_dtype as matmul_matrix_qualifiers_dtype
+from .generic import matrix_qualifiers_dtype as matmul_matrix_qualifiers_dtype  # noqa: E402
 
 __all__ = [
     "advanced",

nvmath/tensor/__init__.py

@@ -2,5 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from ._configuration import *  # noqa: F403
-from .contract import *  # noqa: F403
+from nvmath._utils import module_init_force_cupy_lib_load
+
+module_init_force_cupy_lib_load()
+
+from ._configuration import *  # noqa: E402, F403
+from .contract import *  # noqa: E402, F403

releasenotes/notes/lazy-bindings-init-7a3c9e2b1d4f8c60.yaml

@@ -0,0 +1,19 @@
+---
+# Uncomment a section header and replace with your own note.
+#release_summary:
+#api:
+features:
+  - |
+    ``nvmath`` and ``nvmath.bindings`` now import high-level subpackages and
+    library wrappers on first access instead of loading every wrapper
+    (including optional ``*Mp`` and NVSHMEM modules) during package import.
+    Importing ``nvmath.bindings.cycurand`` therefore no longer pulls in
+    ``fft`` / ``linalg`` / ``sparse`` / ``tensor``. Public names, ``__all__``,
+    and the behavior that optional binding modules are ``None`` when
+    unavailable are unchanged.
+#dependencies:
+#deprecations:
+#fixes:
+#issues:
+#security:
+#doc: