push sheeet

pkgs/development/libraries/science/math/libtorch/bin.nix

@@ -0,0 +1,118 @@
+{
+  callPackage,
+  stdenv,
+  fetchzip,
+  lib,
+  libcxx,
+  llvmPackages,
+  config,
+
+  addDriverRunpath,
+  patchelf,
+  fixDarwinDylibNames,
+
+  cudaSupport ? config.cudaSupport,
+}:
+
+let
+  # The binary libtorch distribution statically links the CUDA
+  # toolkit. This means that we do not need to provide CUDA to
+  # this derivation. However, we should ensure on version bumps
+  # that the CUDA toolkit for `passthru.tests` is still
+  # up-to-date.
+  version = "2.5.0";
+  device = if cudaSupport then "cuda" else "cpu";
+  srcs = import ./binary-hashes.nix version;
+  unavailable = throw "libtorch is not available for this platform";
+in
+stdenv.mkDerivation {
+  inherit version;
+  pname = "libtorch";
+
+  src = fetchzip srcs."${stdenv.hostPlatform.system}-${device}" or unavailable;
+
+  nativeBuildInputs =
+    if stdenv.hostPlatform.isDarwin then
+      [ fixDarwinDylibNames ]
+    else
+      [ patchelf ] ++ lib.optionals cudaSupport [ addDriverRunpath ];
+
+  dontBuild = true;
+  dontConfigure = true;
+  dontStrip = true;
+
+  installPhase = ''
+    # Copy headers and CMake files.
+    mkdir -p $dev
+    cp -r include $dev
+    cp -r share $dev
+
+    install -Dm755 -t $out/lib lib/*${stdenv.hostPlatform.extensions.sharedLibrary}*
+
+    # We do not care about Java support...
+    rm -f $out/lib/lib*jni* 2> /dev/null || true
+
+    # Fix up library paths for split outputs
+    substituteInPlace $dev/share/cmake/Torch/TorchConfig.cmake \
+      --replace \''${TORCH_INSTALL_PREFIX}/lib "$out/lib" \
+
+    substituteInPlace \
+      $dev/share/cmake/Caffe2/Caffe2Targets-release.cmake \
+      --replace \''${_IMPORT_PREFIX}/lib "$out/lib" \
+  '';
+
+  postFixup =
+    let
+      rpath = lib.makeLibraryPath [ stdenv.cc.cc ];
+    in
+    lib.optionalString stdenv.hostPlatform.isLinux ''
+      find $out/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+        echo "setting rpath for $lib..."
+        patchelf --set-rpath "${rpath}:$out/lib" "$lib"
+        ${lib.optionalString cudaSupport ''
+          addDriverRunpath "$lib"
+        ''}
+      done
+    ''
+    + lib.optionalString stdenv.hostPlatform.isDarwin ''
+      for f in $out/lib/*.dylib; do
+          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
+        done
+      done
+      for f in $out/lib/*.dylib; do
+          otool -L $f
+      done
+    '';
+
+  outputs = [
+    "out"
+    "dev"
+  ];
+
+  passthru.tests.cmake = callPackage ./test {
+    inherit cudaSupport;
+  };
+
+  meta = with lib; {
+    description = "C++ API of the PyTorch machine learning framework";
+    homepage = "https://pytorch.org/";
+    sourceProvenance = with sourceTypes; [ binaryNativeCode ];
+    # Includes CUDA and Intel MKL, but redistributions of the binary are not limited.
+    # https://docs.nvidia.com/cuda/eula/index.html
+    # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
+    license = licenses.bsd3;
+    maintainers = with maintainers; [ junjihashimoto ];
+    platforms = [
+      "aarch64-darwin"
+      "x86_64-linux"
+    ];
+  };
+}

pkgs/development/libraries/science/math/libtorch/binary-hashes.nix

@@ -0,0 +1,20 @@
+version:
+builtins.getAttr version {
+  "2.5.0" = {
+    aarch64-darwin-cpu = {
+      name = "libtorch-macos-arm64-2.5.0.zip";
+      url = "https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-2.5.0.zip";
+      hash = "sha256-4d9YKUuvAESBeG/WCUsQfEHwdB2z34grnlwWslj4970=";
+    };
+    x86_64-linux-cpu = {
+      name = "libtorch-cxx11-abi-shared-with-deps-2.5.0-cpu.zip";
+      url = "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcpu.zip";
+      hash = "sha256-gUzPhc4Z8rTPhIm89pPoLP0Ww17ono+/xgMW46E/Tro=";
+    };
+    x86_64-linux-cuda = {
+      name = "libtorch-cxx11-abi-shared-with-deps-2.5.0-cu124.zip";
+      url = "https://download.pytorch.org/libtorch/cu124/libtorch-cxx11-abi-shared-with-deps-2.5.0%2Bcu124.zip";
+      hash = "sha256-UaX47GAwyZ6UmzgY85TeAHmy3u52pBHhiyM5NAz7ens=";
+    };
+  };
+}

pkgs/development/libraries/science/math/libtorch/prefetch.sh

@@ -0,0 +1,38 @@
+#!/usr/bin/env nix-shell
+#!nix-shell -i bash -p nix-prefetch-scripts
+
+set -eou pipefail
+
+version=$1
+
+bucket="https://download.pytorch.org/libtorch"
+CUDA_VERSION=cu124
+
+url_and_key_list=(
+  "aarch64-darwin-cpu $bucket/cpu/libtorch-macos-arm64-${version}.zip libtorch-macos-arm64-${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"
+)
+
+hashfile="binary-hashes-$version.nix"
+echo "  \"$version\" = {" >> $hashfile
+
+for url_and_key in "${url_and_key_list[@]}"; do
+  key=$(echo "$url_and_key" | cut -d' ' -f1)
+  url=$(echo "$url_and_key" | cut -d' ' -f2)
+  name=$(echo "$url_and_key" | cut -d' ' -f3)
+
+  echo "prefetching ${url}..."
+  hash=$(nix --extra-experimental-features nix-command hash to-sri --type sha256 $(nix-prefetch-url --unpack "$url" --name "$name"))
+
+  echo "    $key = {" >> $hashfile
+  echo "      name = \"$name\";" >> $hashfile
+  echo "      url = \"$url\";" >> $hashfile
+  echo "      hash = \"$hash\";" >> $hashfile
+  echo "    };" >> $hashfile
+
+  echo
+done
+
+echo "  };" >> $hashfile
+echo "done."

pkgs/development/libraries/science/math/libtorch/test/CMakeLists.txt

@@ -0,0 +1,5 @@
+cmake_minimum_required(VERSION 3.0)
+find_package(Torch REQUIRED)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+add_executable(test test.cpp)
+target_link_libraries(test "${TORCH_LIBRARIES}")

pkgs/development/libraries/science/math/libtorch/test/default.nix

@@ -0,0 +1,64 @@
+{
+  lib,
+  stdenv,
+  cmake,
+  libtorch-bin,
+  linkFarm,
+  symlinkJoin,
+
+  cudaSupport,
+  cudaPackages ? { },
+}:
+let
+  inherit (cudaPackages) cudatoolkit cudnn;
+
+  cudatoolkit_joined = symlinkJoin {
+    name = "${cudatoolkit.name}-unsplit";
+    paths = [
+      cudatoolkit.out
+      cudatoolkit.lib
+    ];
+  };
+
+  # We do not have access to /run/opengl-driver/lib in the sandbox,
+  # so use a stub instead.
+  cudaStub = linkFarm "cuda-stub" [
+    {
+      name = "libcuda.so.1";
+      path = "${cudatoolkit}/lib/stubs/libcuda.so";
+    }
+  ];
+
+in
+stdenv.mkDerivation {
+  pname = "libtorch-test";
+  version = libtorch-bin.version;
+
+  src = lib.fileset.toSource {
+    root = ./.;
+    fileset = lib.fileset.unions [
+      ./CMakeLists.txt
+      ./test.cpp
+    ];
+  };
+
+  nativeBuildInputs = [ cmake ];
+
+  buildInputs = [ libtorch-bin ] ++ lib.optionals cudaSupport [ cudnn ];
+
+  cmakeFlags = lib.optionals cudaSupport [ "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit_joined}" ];
+
+  doCheck = true;
+
+  installPhase = ''
+    touch $out
+  '';
+
+  checkPhase =
+    lib.optionalString cudaSupport ''
+      LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
+    ''
+    + ''
+      ./test
+    '';
+}

pkgs/development/libraries/science/math/libtorch/test/test.cpp

@@ -0,0 +1,20 @@
+#undef NDEBUG
+#include <cassert>
+
+#include <iostream>
+
+#include <torch/torch.h>
+
+int main() {
+  torch::Tensor tensor = torch::eye(3);
+
+  float checkData[] = {
+    1, 0, 0,
+    0, 1, 0,
+    0, 0, 1
+  };
+
+  torch::Tensor check = torch::from_blob(checkData, {3, 3});
+
+  assert(tensor.allclose(check));
+}