pkgs/development/ocaml-modules/torch/default.nix

   1 { lib
   2 , stdenv
   3 , buildDunePackage
   4 , fetchFromGitHub
   5 , fetchpatch
   6 , cmdliner
   7 , ctypes
   8 , dune-configurator
   9 , npy
  10 , ocaml-compiler-libs
  11 , ppx_custom_printf
  12 , ppx_expect
  13 , ppx_sexp_conv
  14 , sexplib
  15 , stdio
  16 , torch
  17 }:
  18
  19 buildDunePackage rec {
  20   pname = "torch";
  21   version = "0.17";
  22
  23   minimalOCamlVersion = "4.08";
  24
  25   src = fetchFromGitHub {
  26     owner = "LaurentMazare";
  27     repo = "ocaml-${pname}";
  28     rev = version;
  29     hash = "sha256-z/9NUBjeFWE63Z/e8OyzDiy8hrn6qzjaiBH8G9MPeos=";
  30   };
  31
  32   patches = [
  33     # Pytorch 2.0 support. Drop when it reaches a release
  34     (fetchpatch {
  35       url = "https://github.com/LaurentMazare/ocaml-torch/commit/ef7ef30cafecb09e45ec1ed8ce4bedae5947cfa5.patch";
  36       hash = "sha256-smdwKy40iIISp/25L2J4az6KmqFS1soeChBElUyhl5A=";
  37     })
  38   ];
  39
  40   buildInputs = [ dune-configurator ];
  41
  42   propagatedBuildInputs = [
  43     cmdliner
  44     ctypes
  45     npy
  46     ocaml-compiler-libs
  47     ppx_custom_printf
  48     ppx_expect
  49     ppx_sexp_conv
  50     sexplib
  51     stdio
  52     torch
  53     torch.dev
  54   ];
  55
  56   preBuild = "export LIBTORCH=${torch.dev}/";
  57
  58   doCheck = !stdenv.isAarch64;
  59
  60   meta = with lib; {
  61     inherit (src.meta) homepage;
  62     description = "Ocaml bindings to Pytorch";
  63     maintainers = [ maintainers.bcdarwin ];
  64     license = licenses.asl20;
  65   };
  66 }