pkgs/development/python-modules/pytorch/bin.nix

   1 { lib, stdenv
   2 , buildPythonPackage
   3 , fetchurl
   4 , isPy37
   5 , isPy38
   6 , isPy39
   7 , python
   8 , addOpenGLRunpath
   9 , future
  10 , numpy
  11 , patchelf
  12 , pyyaml
  13 , requests
  14 , typing-extensions
  15 }:
  16
  17 let
  18   pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
  19   srcs = import ./binary-hashes.nix version;
  20   unsupported = throw "Unsupported system";
  21   version = "1.8.0";
  22 in buildPythonPackage {
  23   inherit version;
  24
  25   pname = "pytorch";
  26   # Don't forget to update pytorch to the same version.
  27
  28   format = "wheel";
  29
  30   disabled = !(isPy37 || isPy38 || isPy39);
  31
  32   src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
  33
  34   nativeBuildInputs = [
  35     addOpenGLRunpath
  36     patchelf
  37   ];
  38
  39   propagatedBuildInputs = [
  40     future
  41     numpy
  42     pyyaml
  43     requests
  44     typing-extensions
  45   ];
  46
  47   postInstall = ''
  48     # ONNX conversion
  49     rm -rf $out/bin
  50   '';
  51
  52   postFixup = let
  53     rpath = lib.makeLibraryPath [ stdenv.cc.cc.lib ];
  54   in ''
  55     find $out/${python.sitePackages}/torch/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
  56       echo "setting rpath for $lib..."
  57       patchelf --set-rpath "${rpath}:$out/${python.sitePackages}/torch/lib" "$lib"
  58       addOpenGLRunpath "$lib"
  59     done
  60   '';
  61
  62   pythonImportsCheck = [ "torch" ];
  63
  64   meta = with lib; {
  65     description = "Open source, prototype-to-production deep learning platform";
  66     homepage = "https://pytorch.org/";
  67     license = licenses.unfree; # Includes CUDA and Intel MKL.
  68     platforms = platforms.linux;
  69     maintainers = with maintainers; [ danieldk ];
  70   };
  71 }