pkgs/development/python-modules/pytorch-pfn-extras/default.nix

   1 { buildPythonPackage
   2 , fetchFromGitHub
   3 , fetchpatch
   4 , lib
   5 , numpy
   6 , onnx
   7 , packaging
   8 , pytestCheckHook
   9 , pythonAtLeast
  10 , setuptools
  11 , stdenv
  12 , torch
  13 , torchvision
  14 , typing-extensions
  15 }:
  16
  17 buildPythonPackage rec {
  18   pname = "pytorch-pfn-extras";
  19   version = "0.7.6";
  20   pyproject = true;
  21
  22   src = fetchFromGitHub {
  23     owner = "pfnet";
  24     repo = "pytorch-pfn-extras";
  25     rev = "refs/tags/v${version}";
  26     hash = "sha256-vSon/0GxQfaRtSPsQbYAvE3s/F0HEN59VpzE3w1PnVE=";
  27   };
  28
  29   patches = [
  30     (fetchpatch {
  31       name = "relax-setuptools.patch";
  32       url = "https://github.com/pfnet/pytorch-pfn-extras/commit/96abe38c4baa6a144d604bdd4744c55627e55440.patch";
  33       hash = "sha256-85UDGcgJyQS5gINbgpNM58b3XJGvf+ArtGhwJ5EXdhk=";
  34     })
  35   ];
  36
  37   build-system = [
  38     setuptools
  39   ];
  40
  41   dependencies = [ numpy packaging torch typing-extensions ];
  42
  43   nativeCheckInputs = [ onnx pytestCheckHook torchvision ];
  44
  45   pytestFlagsArray = [
  46     # Requires CUDA access which is not possible in the nix environment.
  47     "-m 'not gpu and not mpi'"
  48     "-Wignore::DeprecationWarning"
  49   ];
  50
  51   pythonImportsCheck = [ "pytorch_pfn_extras" ];
  52
  53   disabledTestPaths = [
  54     # Requires optuna which is currently (2022-02-16) marked as broken.
  55     "tests/pytorch_pfn_extras_tests/test_config_types.py"
  56
  57     # requires onnxruntime which was removed because of poor maintainability
  58     # See https://github.com/NixOS/nixpkgs/pull/105951 https://github.com/NixOS/nixpkgs/pull/155058
  59     "tests/pytorch_pfn_extras_tests/onnx_tests/test_annotate.py"
  60     "tests/pytorch_pfn_extras_tests/onnx_tests/test_as_output.py"
  61     "tests/pytorch_pfn_extras_tests/onnx_tests/test_export.py"
  62     "tests/pytorch_pfn_extras_tests/onnx_tests/test_export_testcase.py"
  63     "tests/pytorch_pfn_extras_tests/onnx_tests/test_lax.py"
  64     "tests/pytorch_pfn_extras_tests/onnx_tests/test_load_model.py"
  65     "tests/pytorch_pfn_extras_tests/onnx_tests/test_torchvision.py"
  66     "tests/pytorch_pfn_extras_tests/onnx_tests/utils.py"
  67
  68     # RuntimeError: No Op registered for Gradient with domain_version of 9
  69     "tests/pytorch_pfn_extras_tests/onnx_tests/test_grad.py"
  70   ] ++ lib.optionals (pythonAtLeast "3.12") [
  71     # RuntimeError: Dynamo is not supported on Python 3.12+
  72     "tests/pytorch_pfn_extras_tests/dynamo_tests/test_compile.py"
  73     "tests/pytorch_pfn_extras_tests/test_ops/test_register.py"
  74   ] ++ lib.optionals (stdenv.hostPlatform.isDarwin) [
  75     # torch.distributed is not available on darwin
  76     "tests/pytorch_pfn_extras_tests/training_tests/extensions_tests/test_sharded_snapshot.py"
  77   ] ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
  78     # RuntimeError: internal error
  79     # convolution (e.g. F.conv3d) causes runtime error
  80     "tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy_conv.py"
  81   ];
  82
  83   meta = with lib; {
  84     description = "Supplementary components to accelerate research and development in PyTorch";
  85     homepage = "https://github.com/pfnet/pytorch-pfn-extras";
  86     license = licenses.mit;
  87     maintainers = with maintainers; [ samuela ];
  88   };
  89 }