Copy over Torch from hf-nix and fix the AArch64 build

overlay.nix

@@ -50,6 +50,26 @@ final: prev: {
         });
 
         pyclibrary = python-self.callPackage ./pkgs/python-modules/pyclibrary { };
+
+        mkTorch = callPackage ./pkgs/python-modules/torch/binary { };
+
+        torch-bin_2_8 = mkTorch {
+          version = "2.8";
+          xpuPackages = final.xpuPackages_2025_1;
+        };
+
+        torch-bin_2_9 = mkTorch {
+          version = "2.9";
+          xpuPackages = final.xpuPackages_2025_2;
+        };
+
+        torch_2_8 = callPackage ./pkgs/python-modules/torch/source/2_8 {
+          xpuPackages = final.xpuPackages_2025_1;
+        };
+
+        torch_2_9 = callPackage ./pkgs/python-modules/torch/source/2_9 {
+          xpuPackages = final.xpuPackages_2025_2;
+        };
       }
     )
   ];

pkgs/python-modules/torch/archs.nix

@@ -0,0 +1,128 @@
+{
+  "2.8" = {
+    # https://github.com/pytorch/pytorch/blob/release/2.8/.ci/manywheel/build_cuda.sh
+    capsPerCudaVersion = {
+      "12.9" = [
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      "12.8" = [
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      "12.6" = [
+        "5.0"
+        "6.0"
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+      ];
+      # Not a supported upstream configuration, but keep it around for
+      # builds that fail on newer CUDA versions.
+      "12.4" = [
+        "5.0"
+        "6.0"
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+      ];
+    };
+    # https://github.com/pytorch/pytorch/blob/ba56102387ef21a3b04b357e5b183d48f0afefc7/.ci/docker/manywheel/build.sh#L82
+    supportedTorchRocmArchs = [
+      "gfx900"
+      "gfx906"
+      "gfx908"
+      "gfx90a"
+      "gfx942"
+      "gfx1030"
+      "gfx1100"
+      "gfx1101"
+      "gfx1102"
+      "gfx1200"
+      "gfx1201"
+    ];
+  };
+
+  "2.9" = {
+    # https://github.com/pytorch/pytorch/blob/release/2.9/.ci/manywheel/build_cuda.sh
+    capsPerCudaVersion = {
+      "13.0" = [
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      # NOTE: 12.9 does not seem to be in RC builds, check if needed for final release.
+      #       https://download.pytorch.org/whl/test/torch/
+      "12.9" = [
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      "12.8" = [
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+        "10.0"
+        "12.0"
+      ];
+      "12.6" = [
+        "5.0"
+        "6.0"
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+      ];
+      # Not a supported upstream configuration, but keep it around for
+      # builds that fail on newer CUDA versions.
+      "12.4" = [
+        "5.0"
+        "6.0"
+        "7.0"
+        "7.5"
+        "8.0"
+        "8.6"
+        "9.0"
+      ];
+    };
+
+    supportedTorchRocmArchs = [
+      # https://github.com/pytorch/pytorch/blob/21fec65781bebe867faf209f89bb687ffd236ca4/.ci/docker/manywheel/build.sh#L92
+      "gfx900"
+      "gfx906"
+      "gfx908"
+      "gfx90a"
+      "gfx942"
+      "gfx1030"
+      "gfx1100"
+      "gfx1101"
+      "gfx1102"
+      "gfx1200"
+      "gfx1201"
+    ];
+  };
+}

pkgs/python-modules/torch/binary/default.nix

@@ -0,0 +1,43 @@
+{
+  config,
+  lib,
+  stdenv,
+
+  cudaSupport ? config.cudaSupport,
+  rocmSupport ? config.rocmSupport,
+  xpuSupport ? (config.xpuSupport or false),
+
+  callPackage,
+  cudaPackages,
+  rocmPackages,
+}:
+
+{
+  xpuPackages,
+  version,
+}:
+
+let
+  system = stdenv.hostPlatform.system;
+  flattenVersion = version: lib.replaceStrings [ "." ] [ "" ] (lib.versions.pad 2 version);
+  framework =
+    if cudaSupport then
+      "cu${flattenVersion cudaPackages.cudaMajorMinorVersion}"
+    else if rocmSupport then
+      "rocm${flattenVersion (lib.versions.majorMinor rocmPackages.rocm.version)}"
+    else if xpuSupport then
+      "xpu"
+    else
+      "cpu";
+  torchVersions = builtins.fromJSON (builtins.readFile ./torch-versions-hash.json);
+  torchBySystem = torchVersions.${version} or (throw "Unsupported torch version: ${version}");
+  torchByFramework =
+    torchBySystem.${system} or (throw "Unsupported system: ${system} for torch version: ${version}");
+  urlHash =
+    torchByFramework.${framework}
+      or (throw "Unsupported framework: ${framework} for torch version: ${version} on system: ${system}");
+in
+callPackage ./generic.nix {
+  inherit xpuPackages;
+  inherit (urlHash) url hash version;
+}