From 33b21aec60575e499244b9f34c6055b107e7e39d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Sun, 30 Aug 2020 11:34:32 +0200
Subject: [PATCH] pytorch-bin: init at 1.6.0

---
 .../python-modules/pytorch/bin.nix            | 68 +++++++++++++++++++
 .../python-modules/pytorch/binary-hashes.nix  | 10 +++
 .../python-modules/pytorch/default.nix        |  4 +-
 pkgs/top-level/python-packages.nix            |  4 ++
 4 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 pkgs/development/python-modules/pytorch/bin.nix
 create mode 100644 pkgs/development/python-modules/pytorch/binary-hashes.nix

diff --git a/pkgs/development/python-modules/pytorch/bin.nix b/pkgs/development/python-modules/pytorch/bin.nix
new file mode 100644
index 000000000000..b1b662e95b3e
--- /dev/null
+++ b/pkgs/development/python-modules/pytorch/bin.nix
@@ -0,0 +1,68 @@
+{ stdenv
+, buildPythonPackage
+, fetchurl
+, isPy37
+, isPy38
+, python
+, nvidia_x11
+, addOpenGLRunpath
+, future
+, numpy
+, patchelf
+, pyyaml
+, requests
+}:
+
+let
+  pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
+  platform = if stdenv.isDarwin then "darwin" else "linux";
+  srcs = import ./binary-hashes.nix;
+  unsupported = throw "Unsupported system";
+in buildPythonPackage {
+  pname = "pytorch";
+  # Don't forget to update pytorch to the same version.
+  version = "1.6.0";
+
+  format = "wheel";
+
+  disabled = !(isPy37 || isPy38);
+
+  src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
+
+  nativeBuildInputs = [
+    addOpenGLRunpath
+    patchelf
+  ];
+
+  propagatedBuildInputs = [
+    future
+    numpy
+    pyyaml
+    requests
+  ];
+
+  postInstall = ''
+    # ONNX conversion
+    rm -rf $out/bin
+  '';
+
+  postFixup = let
+    rpath = stdenv.lib.makeLibraryPath [ stdenv.cc.cc.lib nvidia_x11 ];
+  in ''
+    find $out/${python.sitePackages}/torch/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      echo "setting rpath for $lib..."
+      patchelf --set-rpath "${rpath}:$out/${python.sitePackages}/torch/lib" "$lib"
+      addOpenGLRunpath "$lib"
+    done
+  '';
+
+  pythonImportsCheck = [ "torch" ];
+
+  meta = with stdenv.lib; {
+    description = "Open source, prototype-to-production deep learning platform";
+    homepage = "https://pytorch.org/";
+    license = licenses.unfree; # Includes CUDA and Intel MKL.
+    platforms = platforms.linux;
+    maintainers = with maintainers; [ danieldk ];
+  };
+}
diff --git a/pkgs/development/python-modules/pytorch/binary-hashes.nix b/pkgs/development/python-modules/pytorch/binary-hashes.nix
new file mode 100644
index 000000000000..481dc6825c7d
--- /dev/null
+++ b/pkgs/development/python-modules/pytorch/binary-hashes.nix
@@ -0,0 +1,10 @@
+{
+  x86_64-linux-37 = {
+    url = "https://download.pytorch.org/whl/cu102/torch-1.6.0-cp37-cp37m-linux_x86_64.whl";
+    sha256 = "0xhwv68j8gvahfzcp43bqp2x71iwv6zjhkw2f1hb82xps40mrml7";
+  };
+  x86_64-linux-38 = {
+    url = "https://download.pytorch.org/whl/cu102/torch-1.6.0-cp38-cp38-linux_x86_64.whl";
+    sha256 = "05m2l04wqzw5xvjam6zwvlmc3979cksl3hrdqc2aikrv4hz8fmsk";
+  };
+}
diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix
index a767a913abbb..9bdead706d7b 100644
--- a/pkgs/development/python-modules/pytorch/default.nix
+++ b/pkgs/development/python-modules/pytorch/default.nix
@@ -104,8 +104,10 @@ let
     "LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH ";
 
 in buildPythonPackage rec {
-  version = "1.6.0";
   pname = "pytorch";
+  # Don't forget to update pytorch-bin to the same version.
+  version = "1.6.0";
+
   disabled = !isPy3k;
 
   outputs = [
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index cfed029af890..72fd49c64f78 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -4100,6 +4100,10 @@ in {
     cudaSupport = pkgs.config.cudaSupport or false;
   };
 
+  pytorch-bin = callPackage ../development/python-modules/pytorch/bin.nix {
+    inherit (pkgs.linuxPackages) nvidia_x11;
+  };
+
   pyro-ppl = callPackage ../development/python-modules/pyro-ppl {};
 
   opt-einsum = if isPy27 then