|
|
|
@ -8,19 +8,19 @@ |
|
|
|
|
, numpy |
|
|
|
|
, pytest-xdist |
|
|
|
|
, pytestCheckHook |
|
|
|
|
, tensorflow |
|
|
|
|
, tensorflow-datasets |
|
|
|
|
}: |
|
|
|
|
|
|
|
|
|
buildPythonPackage rec { |
|
|
|
|
pname = "optax"; |
|
|
|
|
# As of 2022-01-06, the latest stable version (0.1.0) has broken tests that are fixed |
|
|
|
|
# in https://github.com/deepmind/optax/commit/d6633365d84eb6f2c0df0c52b630481a349ce562 |
|
|
|
|
version = "unstable-2022-01-05"; |
|
|
|
|
version = "0.1.1"; |
|
|
|
|
|
|
|
|
|
src = fetchFromGitHub { |
|
|
|
|
owner = "deepmind"; |
|
|
|
|
repo = pname; |
|
|
|
|
rev = "5ec5541b3486224b22e950480ff639ceaf5098f7"; |
|
|
|
|
sha256 = "1q8cxc42a5xais2ll1l238cnn3l7w28savhgiz0lg01ilz2ysbli"; |
|
|
|
|
rev = "v${version}"; |
|
|
|
|
hash = "sha256-s/BcqzhdfWzR61MStusUPQtuT4+t8NcC5gBGiGggFqw="; |
|
|
|
|
}; |
|
|
|
|
|
|
|
|
|
buildInputs = [ jaxlib ]; |
|
|
|
@ -35,6 +35,8 @@ buildPythonPackage rec { |
|
|
|
|
dm-haiku |
|
|
|
|
pytest-xdist |
|
|
|
|
pytestCheckHook |
|
|
|
|
tensorflow |
|
|
|
|
tensorflow-datasets |
|
|
|
|
]; |
|
|
|
|
pytestFlagsArray = [ "-n $NIX_BUILD_CORES" ]; |
|
|
|
|
|
|
|
|
@ -45,8 +47,7 @@ buildPythonPackage rec { |
|
|
|
|
disabledTestPaths = [ |
|
|
|
|
# Requires `flax` which depends on `optax` creating circular dependency. |
|
|
|
|
"optax/_src/equivalence_test.py" |
|
|
|
|
# Require `tensorflow_datasets` which isn't packaged in `nixpkgs`. |
|
|
|
|
"examples/datasets_test.py" |
|
|
|
|
# See https://github.com/deepmind/optax/issues/323. |
|
|
|
|
"examples/lookahead_mnist_test.py" |
|
|
|
|
]; |
|
|
|
|
|
|
|
|
|