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

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