From 460e7105cd2b7c35d66c1794e9281743e955166a Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Fri, 10 Apr 2026 07:55:09 +0000
Subject: [PATCH] python3Packages.transformer-engine: 2.12 -> 2.14

Diff: https://github.com/NVIDIA/TransformerEngine/compare/v2.12...v2.14

Changelog:
https://github.com/NVIDIA/TransformerEngine/releases/tag/v2.14
---
 .../transformer-engine/cuda-libs-paths.patch     | 16 +++++++++-------
 .../transformer-engine/default.nix               | 13 +++++++------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/pkgs/development/python-modules/transformer-engine/cuda-libs-paths.patch b/pkgs/development/python-modules/transformer-engine/cuda-libs-paths.patch
index bdf80680acbb..49ddd0553f49 100644
--- a/pkgs/development/python-modules/transformer-engine/cuda-libs-paths.patch
+++ b/pkgs/development/python-modules/transformer-engine/cuda-libs-paths.patch
@@ -1,8 +1,8 @@
 diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
-index 02388d2e..f2eb337c 100644
+index 40933f17..f2eb337c 100644
 --- a/transformer_engine/common/__init__.py
 +++ b/transformer_engine/common/__init__.py
-@@ -239,117 +239,7 @@ def _get_sys_extension() -> str:
+@@ -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."""
  
@@ -12,11 +12,13 @@ index 02388d2e..f2eb337c 100644
 -        return ""
 -
 -    # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia"
--    # above doesn't through. However, they don't set "__file__" attribute.
--    if nvidia.__file__ is None:
--        return ""
+-    # 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 = Path(nvidia.__file__).parent / "cuda_runtime"
+-    include_dir = nvidia_root / "cuda_runtime"
 -    return str(include_dir) if include_dir.exists() else ""
 -
 -
@@ -121,7 +123,7 @@ index 02388d2e..f2eb337c 100644
  
  
  @functools.lru_cache(maxsize=None)
-@@ -364,18 +254,9 @@ if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE
+@@ -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.
diff --git a/pkgs/development/python-modules/transformer-engine/default.nix b/pkgs/development/python-modules/transformer-engine/default.nix
index 9a29b6b75e4b..0e9fb9c59588 100644
--- a/pkgs/development/python-modules/transformer-engine/default.nix
+++ b/pkgs/development/python-modules/transformer-engine/default.nix
@@ -77,8 +77,9 @@ let
 in
 buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
   pname = "transformer-engine";
-  version = "2.12";
+  version = "2.14";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "NVIDIA";
@@ -86,7 +87,7 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
     tag = "v${finalAttrs.version}";
     # Their CMakeLists.txt does not easily let us inject dependencies
     fetchSubmodules = true;
-    hash = "sha256-/e11kacSYPKdjVEKAo3x/CarzKhO3tiTsMjYWLzHbls=";
+    hash = "sha256-yxcUn75blB5ssEqGXZFDUrBv2/WM8yzumJCt5olV5Po=";
   };
 
   patches =
@@ -102,8 +103,8 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
       # https://github.com/NVIDIA/TransformerEngine/pull/2832
       (fetchpatch {
         name = "fix-cuda-arch-cmake-logic";
-        url = "https://github.com/GaetanLepage/TransformerEngine/commit/a3cf63e0d03dd9af1d494854949387f1ae677bf0.patch";
-        hash = "sha256-g2aIF0fROsExEjuNiyI62/rrCOXYyOjyQIOn6rCrUyI=";
+        url = "https://github.com/NVIDIA/TransformerEngine/commit/fca261ecd09c318d22e7eeebda79632eed8cb9e4.patch";
+        hash = "sha256-nph01cIfmjN7RUFZif5ORoz29CEHWwetiHMEZVnnOyY=";
       })
     ]
     ++ optionals withNvshmem [
@@ -118,8 +119,8 @@ buildPythonPackage.override { stdenv = backendStdenv; } (finalAttrs: {
       # https://github.com/NVIDIA/TransformerEngine/pull/2835
       (fetchpatch {
         name = "fix-jax-extension-build-with-mpi";
-        url = "https://github.com/GaetanLepage/TransformerEngine/commit/f68cd3cab34972a899ad0069e2c4ee806e8bc6fb.patch";
-        hash = "sha256-u0ljg1FwY0QjR+ETswpzWV+Sbv00JHI5CSrNQ/9zsuA=";
+        url = "https://github.com/NVIDIA/TransformerEngine/commit/2dd31bb849e83cce51c7d169db883862063d3a95.patch";
+        hash = "sha256-QSRMetseYPGGZCgGkS9rIj9nJdazCD4hv2IgPc+ClSM=";
       })
     ];