|
|
|
@ -1,14 +1,14 @@ |
|
|
|
|
{ buildPythonPackage |
|
|
|
|
{ lib |
|
|
|
|
, buildPythonPackage |
|
|
|
|
, cloudpickle |
|
|
|
|
, deepdish |
|
|
|
|
, deepmerge |
|
|
|
|
, dm-haiku |
|
|
|
|
, fetchFromGitHub |
|
|
|
|
, jaxlib |
|
|
|
|
, lib |
|
|
|
|
, poetry |
|
|
|
|
, pytestCheckHook |
|
|
|
|
, torch |
|
|
|
|
, pythonOlder |
|
|
|
|
, pyyaml |
|
|
|
|
, sh |
|
|
|
|
, tables |
|
|
|
@ -16,6 +16,7 @@ |
|
|
|
|
, tensorboardx |
|
|
|
|
, tensorflow |
|
|
|
|
, toolz |
|
|
|
|
, torch |
|
|
|
|
, treex |
|
|
|
|
, typing-extensions |
|
|
|
|
}: |
|
|
|
@ -25,6 +26,8 @@ buildPythonPackage rec { |
|
|
|
|
version = "0.8.6"; |
|
|
|
|
format = "pyproject"; |
|
|
|
|
|
|
|
|
|
disabled = pythonOlder "3.7"; |
|
|
|
|
|
|
|
|
|
src = fetchFromGitHub { |
|
|
|
|
owner = "poets-ai"; |
|
|
|
|
repo = pname; |
|
|
|
@ -34,8 +37,10 @@ buildPythonPackage rec { |
|
|
|
|
|
|
|
|
|
# The cloudpickle constraint is too strict. wandb is marked as an optional |
|
|
|
|
# dependency but `buildPythonPackage` doesn't seem to respect that setting. |
|
|
|
|
# Python constraint: https://github.com/poets-ai/elegy/issues/244 |
|
|
|
|
postPatch = '' |
|
|
|
|
substituteInPlace pyproject.toml \ |
|
|
|
|
--replace 'python = ">=3.7,<3.10"' 'python = ">=3.7"' \ |
|
|
|
|
--replace 'cloudpickle = "^1.5.0"' 'cloudpickle = "*"' \ |
|
|
|
|
--replace 'wandb = { version = "^0.12.10", optional = true }' "" |
|
|
|
|
''; |
|
|
|
@ -44,7 +49,9 @@ buildPythonPackage rec { |
|
|
|
|
poetry |
|
|
|
|
]; |
|
|
|
|
|
|
|
|
|
buildInputs = [ jaxlib ]; |
|
|
|
|
buildInputs = [ |
|
|
|
|
jaxlib |
|
|
|
|
]; |
|
|
|
|
|
|
|
|
|
propagatedBuildInputs = [ |
|
|
|
|
cloudpickle |
|
|
|
@ -75,6 +82,8 @@ buildPythonPackage rec { |
|
|
|
|
# Fails with `Could not find compiler for platform Host: NOT_FOUND: could not find registered compiler for platform Host -- check target linkage`. |
|
|
|
|
# Runs fine in docker with Ubuntu 22.04. I suspect the issue is the sandboxing in `nixpkgs` but not sure. |
|
|
|
|
"test_saved_model_poly" |
|
|
|
|
# AttributeError: module 'jax' has no attribute 'tree_multimap' |
|
|
|
|
"DataLoaderTestCase" |
|
|
|
|
]; |
|
|
|
|
|
|
|
|
|
meta = with lib; { |
|
|
|
|