nixpkgs/pkgs/development/python-modules/transformer-engine/cuda-libs-paths.patch

diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
index 40933f17..f2eb337c 100644
--- a/transformer_engine/common/__init__.py
+++ b/transformer_engine/common/__init__.py
@@ -239,119 +239,7 @@ def _get_sys_extension() -> str:
 def _nvidia_cudart_include_dir() -> str:
     """Returns the include directory for cuda_runtime.h if exists in python environment."""
 
-    try:
-        import nvidia
-    except ModuleNotFoundError:
-        return ""
-
-    # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia"
-    # above doesn't throw. However, they don't set "__file__" attribute.
-    if nvidia.__file__ is not None:
-        nvidia_root = Path(nvidia.__file__).parent
-    else:
-        nvidia_root = Path(nvidia.__path__[0])  # namespace package
-
-    include_dir = nvidia_root / "cuda_runtime"
-    return str(include_dir) if include_dir.exists() else ""
-
-
-@functools.lru_cache(maxsize=None)
-def _load_cuda_library_from_python(lib_name: str, strict: bool = False):
-    """
-    Attempts to load shared object file installed via python packages.
-
-    `lib_name` : Name of package as found in the `nvidia` dir in python environment.
-    `strict` : If set to `True`, throw an error if lib is not found.
-    """
-
-    ext = _get_sys_extension()
-    nvidia_dir = os.path.join(sysconfig.get_path("purelib"), "nvidia")
-
-    # PyPI packages provided by nvidia libs exist
-    # in 4 possible locations inside `nvidia`.
-    # Check by order of priority.
-    path_found = False
-    if os.path.isdir(os.path.join(nvidia_dir, "cu13", lib_name)):
-        so_paths = glob.glob(os.path.join(nvidia_dir, "cu13", lib_name, f"lib/lib*{ext}.*[0-9]"))
-        path_found = len(so_paths) > 0
-
-    if not path_found and os.path.isdir(os.path.join(nvidia_dir, "cu13")):
-        so_paths = glob.glob(os.path.join(nvidia_dir, "cu13", f"lib/lib{lib_name}*{ext}.*[0-9]"))
-        path_found = len(so_paths) > 0
-
-    if not path_found and os.path.isdir(os.path.join(nvidia_dir, lib_name)):
-        so_paths = glob.glob(os.path.join(nvidia_dir, lib_name, f"lib/lib*{ext}.*[0-9]"))
-        path_found = len(so_paths) > 0
-
-    if not path_found:
-        so_paths = glob.glob(os.path.join(nvidia_dir, f"cuda_{lib_name}", f"lib/lib*{ext}.*[0-9]"))
-        path_found = len(so_paths) > 0
-
-    ctypes_handles = []
-
-    if path_found:
-        for so_path in so_paths:
-            ctypes_handles.append(ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL))
-
-    if strict and not path_found:
-        raise RuntimeError(f"{lib_name} shared object not found.")
-
-    return path_found, ctypes_handles
-
-
-@functools.lru_cache(maxsize=None)
-def _load_cuda_library_from_system(lib_name: str):
-    """
-    Attempts to load shared object file installed via system/cuda-toolkit.
-
-    `lib_name`: Name of library to load without extension or `lib` prefix.
-    """
-
-    # Where to look for the shared lib in decreasing order of preference.
-    paths = (
-        os.environ.get(f"{lib_name.upper()}_HOME"),
-        os.environ.get(f"{lib_name.upper()}_PATH"),
-        os.environ.get("CUDA_HOME"),
-        os.environ.get("CUDA_PATH"),
-        "/usr/local/cuda",
-    )
-
-    for path in paths:
-        if path is None:
-            continue
-        libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True)
-        libs = [lib for lib in libs if "stub" not in lib]
-        libs.sort(reverse=True, key=os.path.basename)
-        if libs:
-            return True, ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL)
-
-    # Search in LD_LIBRARY_PATH.
-    try:
-        _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL)
-        return True, _lib_handle
-    except OSError:
-        return False, None
-
-
-@functools.lru_cache(maxsize=None)
-def _load_cuda_library(lib_name: str):
-    """
-    Load given shared library.
-    Prioritize loading from system/toolkit
-    before checking python packages.
-    """
-
-    # Attempt to locate library in system.
-    found, handle = _load_cuda_library_from_system(lib_name)
-    if found:
-        return True, handle
-
-    # Attempt to locate library in Python dist-packages.
-    found, handle = _load_cuda_library_from_python(lib_name)
-    if found:
-        return False, handle
-
-    raise RuntimeError(f"{lib_name} shared object not found.")
+    return "@cudart_include_dir@"
 
 
 @functools.lru_cache(maxsize=None)
@@ -366,18 +254,9 @@ if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE
     # `_load_cuda_library` is used for packages that must be loaded
     # during runtime. Both system and pypi packages are searched
     # and an error is thrown if not found.
-    _, _CUDNN_LIB_CTYPES = _load_cuda_library("cudnn")
-    system_nvrtc, _NVRTC_LIB_CTYPES = _load_cuda_library("nvrtc")
-    system_curand, _CURAND_LIB_CTYPES = _load_cuda_library("curand")
-
-    # This additional step is necessary to be able to install TE wheels
-    # and import TE (without any guards) in an environment where the cuda
-    # toolkit might be absent without being guarded
-    load_libs_for_no_ctk = not system_nvrtc and not system_curand
-    if load_libs_for_no_ctk:
-        _CUBLAS_LIB_CTYPES = _load_cuda_library_from_python("cublas", strict=True)
-        _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart", strict=True)
-        _CUDNN_ALL_LIB_CTYPES = _load_cuda_library_from_python("cudnn", strict=True)
+    _CUDNN_LIB_CTYPES = ctypes.CDLL("@libcudnn_so@", mode=ctypes.RTLD_GLOBAL)
+    _NVRTC_LIB_CTYPES = ctypes.CDLL("@libnvrtc_so@", mode=ctypes.RTLD_GLOBAL)
+    _CURAND_LIB_CTYPES = ctypes.CDLL("@libcurand_so@", mode=ctypes.RTLD_GLOBAL)
 
     _TE_LIB_CTYPES = _load_core_library()