pkgs/development/python-modules/torchsde/default.nix

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4
   5 # build-system
   6 , setuptools
   7
   8 # dependencies
   9 , boltons
  10 , numpy
  11 , scipy
  12 , torch
  13 , trampoline
  14
  15 # tests
  16 , pytest7CheckHook
  17 }:
  18
  19 buildPythonPackage rec {
  20   pname = "torchsde";
  21   version = "0.2.6";
  22   format = "pyproject";
  23
  24   src = fetchFromGitHub {
  25     owner = "google-research";
  26     repo = "torchsde";
  27     rev = "refs/tags/v${version}";
  28     hash = "sha256-D0p2tL/VvkouXrXfRhMuCq8wMtzeoBTppWEG5vM1qCo=";
  29   };
  30
  31   postPatch = ''
  32     substituteInPlace setup.py \
  33       --replace "numpy==1.19.*" "numpy" \
  34       --replace "scipy==1.5.*" "scipy"
  35   '';
  36
  37   nativeBuildInputs = [
  38     setuptools
  39   ];
  40
  41   propagatedBuildInputs = [
  42     boltons
  43     numpy
  44     scipy
  45     torch
  46     trampoline
  47   ];
  48
  49   pythonImportsCheck = [ "torchsde" ];
  50
  51   nativeCheckInputs = [
  52     pytest7CheckHook
  53   ];
  54
  55   disabledTests = [
  56     # RuntimeError: a view of a leaf Variable that requires grad is being used in an in-place operation.
  57     "test_adjoint"
  58   ];
  59
  60   meta = with lib; {
  61     changelog = "https://github.com/google-research/torchsde/releases/tag/v${version}";
  62     description = "Differentiable SDE solvers with GPU support and efficient sensitivity analysis";
  63     homepage = "https://github.com/google-research/torchsde";
  64     license = licenses.asl20;
  65     maintainers = teams.tts.members;
  66   };
  67 }