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

   1 { lib
   2 , buildPythonPackage
   3 , cloudpickle
   4 , deepdish
   5 , deepmerge
   6 , dm-haiku
   7 , fetchFromGitHub
   8 , fetchpatch
   9 , jaxlib
  10 , poetry-core
  11 , pytestCheckHook
  12 , pythonOlder
  13 , pyyaml
  14 , sh
  15 , tables
  16 , tabulate
  17 , tensorboardx
  18 , tensorflow
  19 , toolz
  20 , torch
  21 , treex
  22 , typing-extensions
  23 }:
  24
  25 buildPythonPackage rec {
  26   pname = "elegy";
  27   version = "0.8.6";
  28   format = "pyproject";
  29
  30   disabled = pythonOlder "3.7";
  31
  32   src = fetchFromGitHub {
  33     owner = "poets-ai";
  34     repo = pname;
  35     rev = "refs/tags/${version}";
  36     hash = "sha256-FZmLriYhsX+zyQKCtCjbOy6MH+AvjzHRNUyaDSXGlLI=";
  37   };
  38
  39   patches = [
  40     (fetchpatch {
  41       name = "use-poetry-core.patch";
  42       url = "https://github.com/poets-ai/elegy/commit/0ed472882f470ed9eb7a63b8a537ffabe7e19aa7.patch";
  43       hash = "sha256-nO/imHo7tEsiZh+64CF/M4eXQ1so3IunVhv8CvYP1ks=";
  44     })
  45   ];
  46
  47   # The cloudpickle constraint is too strict. wandb is marked as an optional
  48   # dependency but `buildPythonPackage` doesn't seem to respect that setting.
  49   # Python constraint: https://github.com/poets-ai/elegy/issues/244
  50   postPatch = ''
  51     substituteInPlace pyproject.toml \
  52       --replace 'python = ">=3.7,<3.10"' 'python = ">=3.7"' \
  53       --replace 'cloudpickle = "^1.5.0"' 'cloudpickle = "*"' \
  54       --replace 'wandb = { version = "^0.12.10", optional = true }' ""
  55   '';
  56
  57   nativeBuildInputs = [
  58     poetry-core
  59   ];
  60
  61   buildInputs = [
  62     jaxlib
  63   ];
  64
  65   propagatedBuildInputs = [
  66     cloudpickle
  67     deepdish
  68     deepmerge
  69     dm-haiku
  70     pyyaml
  71     tables
  72     tabulate
  73     tensorboardx
  74     toolz
  75     treex
  76     typing-extensions
  77   ];
  78
  79   pythonImportsCheck = [
  80     "elegy"
  81   ];
  82
  83   nativeCheckInputs = [
  84     pytestCheckHook
  85     sh
  86     tensorflow
  87     torch
  88   ];
  89
  90   disabledTests = [
  91     # Fails with `Could not find compiler for platform Host: NOT_FOUND: could not find registered compiler for platform Host -- check target linkage`.
  92     # Runs fine in docker with Ubuntu 22.04. I suspect the issue is the sandboxing in `nixpkgs` but not sure.
  93     "test_saved_model_poly"
  94     # AttributeError: module 'jax' has no attribute 'tree_multimap'
  95     "DataLoaderTestCase"
  96   ];
  97
  98   meta = with lib; {
  99     description = "Neural Networks framework based on Jax inspired by Keras and Haiku";
 100     homepage = "https://github.com/poets-ai/elegy";
 101     changelog = "https://github.com/poets-ai/elegy/releases/tag/${version}";
 102     license = licenses.asl20;
 103     maintainers = with maintainers; [ ndl ];
 104   };
 105 }