From 31fc0598df13466e22db21543f09bb68635c0b34 Mon Sep 17 00:00:00 2001
From: Nikita Patsakula <nikita.patsakula@gmail.com>
Date: Wed, 1 May 2024 20:19:15 +0300
Subject: [PATCH] libtorch-bin: 2.0.0 -> 2.3.0

---
 .../libraries/science/math/libtorch/bin.nix   | 13 +++++++---
 .../science/math/libtorch/binary-hashes.nix   | 24 +++++++++----------
 .../science/math/libtorch/prefetch.sh         |  5 ++--
 3 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/pkgs/development/libraries/science/math/libtorch/bin.nix b/pkgs/development/libraries/science/math/libtorch/bin.nix
index 47a6249ebdc6..5b08cfa7dcf0 100644
--- a/pkgs/development/libraries/science/math/libtorch/bin.nix
+++ b/pkgs/development/libraries/science/math/libtorch/bin.nix
@@ -3,6 +3,7 @@
 , fetchzip
 , lib
 , libcxx
+, llvmPackages
 
 , addOpenGLRunpath
 , patchelf
@@ -17,11 +18,11 @@ let
   # this derivation. However, we should ensure on version bumps
   # that the CUDA toolkit for `passthru.tests` is still
   # up-to-date.
-  version = "2.0.0";
+  version = "2.3.0";
   device = if cudaSupport then "cuda" else "cpu";
   srcs = import ./binary-hashes.nix version;
   unavailable = throw "libtorch is not available for this platform";
-  libcxx-for-libtorch = if stdenv.hostPlatform.system == "x86_64-darwin" then libcxx else stdenv.cc.cc.lib;
+  libcxx-for-libtorch = if stdenv.isDarwin then libcxx else stdenv.cc.cc.lib;
 in stdenv.mkDerivation {
   inherit version;
   pname = "libtorch";
@@ -71,6 +72,9 @@ in stdenv.mkDerivation {
         otool -L $f
     done
     for f in $out/lib/*.dylib; do
+      if otool -L $f | grep "@rpath/libomp.dylib" >& /dev/null; then
+        install_name_tool -change "@rpath/libomp.dylib" ${llvmPackages.openmp}/lib/libomp.dylib $f
+      fi
       install_name_tool -id $out/lib/$(basename $f) $f || true
       for rpath in $(otool -L $f | grep rpath | awk '{print $1}');do
         install_name_tool -change $rpath $out/lib/$(basename $rpath) $f
@@ -99,6 +103,9 @@ in stdenv.mkDerivation {
     # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
     license = licenses.bsd3;
     maintainers = with maintainers; [ junjihashimoto ];
-    platforms = platforms.unix;
+    platforms = [
+      "aarch64-darwin"
+      "x86_64-linux"
+    ];
   };
 }
diff --git a/pkgs/development/libraries/science/math/libtorch/binary-hashes.nix b/pkgs/development/libraries/science/math/libtorch/binary-hashes.nix
index 83533d89b1e8..dd8104ad5f0d 100644
--- a/pkgs/development/libraries/science/math/libtorch/binary-hashes.nix
+++ b/pkgs/development/libraries/science/math/libtorch/binary-hashes.nix
@@ -1,19 +1,19 @@
-version : builtins.getAttr version {
-  "2.0.0" = {
-    x86_64-darwin-cpu = {
-      name = "libtorch-macos-2.0.0.zip";
-      url = "https://download.pytorch.org/libtorch/cpu/libtorch-macos-2.0.0.zip";
-      hash = "sha256-u6y5IeYoiOC0yQ/k6JCChDs9lXWccLxUorgR8L62lkM=";
+version: builtins.getAttr version {
+ "2.3.0" = {
+     aarch64-darwin-cpu = {
+      name = "libtorch-macos-arm64-2.3.0.zip";
+      url = "https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.3.0.zip";
+      hash = "sha256-l4qY1jbsECN5qm7dWJ2jSvKuQwQ/HI6G6Vq1Kh2rxkM=";
     };
     x86_64-linux-cpu = {
-      name = "libtorch-cxx11-abi-shared-with-deps-2.0.0-cpu.zip";
-      url = "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcpu.zip";
-      hash = "sha256-BoZQ2MC1CDVVGfX3SHC3mEpLGWO8XK7AcLcHJXDsXuc=";
+      name = "libtorch-cxx11-abi-shared-with-deps-2.3.0-cpu.zip";
+      url = "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.3.0%2Bcpu.zip";
+      hash = "sha256-dKAk6UusK2eQIcP0oMXh9cnufMpy5Ph4SGPkIPPV6ds=";
     };
     x86_64-linux-cuda = {
-      name = "libtorch-cxx11-abi-shared-with-deps-2.0.0-cu118.zip";
-      url = "https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcu118.zip";
-      hash = "sha256-Dpw9kQdA1NI9EOT7JBKwQP4wZT6lizcnKKTQ8WVJCZc=";
+      name = "libtorch-cxx11-abi-shared-with-deps-2.3.0-cu121.zip";
+      url = "https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.3.0%2Bcu121.zip";
+      hash = "sha256-6B+NF6q78I2WKFudn8bK+eNYDi1zQ7mdgv06fZbm2rE=";
     };
   };
 }
diff --git a/pkgs/development/libraries/science/math/libtorch/prefetch.sh b/pkgs/development/libraries/science/math/libtorch/prefetch.sh
index 5c6d60ae8b20..2682fe677976 100755
--- a/pkgs/development/libraries/science/math/libtorch/prefetch.sh
+++ b/pkgs/development/libraries/science/math/libtorch/prefetch.sh
@@ -6,10 +6,11 @@ set -eou pipefail
 version=$1
 
 bucket="https://download.pytorch.org/libtorch"
-CUDA_VERSION=cu116
+CUDA_VERSION=cu121
 
 url_and_key_list=(
-  "x86_64-darwin-cpu $bucket/cpu/libtorch-macos-${version}.zip libtorch-macos-${version}.zip"
+  "aarch64-darwin-cpu $bucket/cpu/libtorch-macos-arm64-${version}.zip libtorch-macos-arm64-${version}.zip"
+  "x86_64-darwin-cpu $bucket/cpu/libtorch-macos-x86_64-${version}.zip libtorch-macos-x86_64-${version}.zip"
   "x86_64-linux-cpu $bucket/cpu/libtorch-cxx11-abi-shared-with-deps-${version}%2Bcpu.zip libtorch-cxx11-abi-shared-with-deps-${version}-cpu.zip"
   "x86_64-linux-cuda $bucket/${CUDA_VERSION}/libtorch-cxx11-abi-shared-with-deps-${version}%2B${CUDA_VERSION}.zip libtorch-cxx11-abi-shared-with-deps-${version}-${CUDA_VERSION}.zip"
 )