pyrox.dev
lol
python310Packages.optax: 0.1.1 -> 0.1.3
authored by
Jonas Heinrich
and committed by
Yt
dcb0d8e9
aaaedf2c
+18
-20
pkgs/development/python-modules/optax/default.nix
+18
-20
pkgs/development/python-modules/optax/default.nix
···
1
1
{ absl-py
2
2
, buildPythonPackage
3
3
, chex
4
-
, dm-haiku
5
4
, fetchFromGitHub
6
5
, jaxlib
7
6
, lib
8
7
, numpy
9
-
, pytest-xdist
10
-
, pytestCheckHook
11
-
, tensorflow
12
-
, tensorflow-datasets
8
+
, callPackage
13
9
}:
14
10
15
11
buildPythonPackage rec {
16
12
pname = "optax";
17
-
version = "0.1.1";
13
+
version = "0.1.3";
18
14
19
15
src = fetchFromGitHub {
20
16
owner = "deepmind";
21
17
repo = pname;
22
18
rev = "v${version}";
23
-
hash = "sha256-s/BcqzhdfWzR61MStusUPQtuT4+t8NcC5gBGiGggFqw=";
19
+
hash = "sha256-XAYztMBQpLBHNuNED/iodbwIMJSN/0GxdmTGQ5jD9Ws=";
24
20
};
25
21
22
+
outputs = [
23
+
"out"
24
+
"testsout"
25
+
];
26
+
26
27
buildInputs = [ jaxlib ];
27
28
28
29
propagatedBuildInputs = [
···
31
32
numpy
32
33
];
33
34
34
-
checkInputs = [
35
-
dm-haiku
36
-
pytest-xdist
37
-
pytestCheckHook
38
-
tensorflow
39
-
tensorflow-datasets
40
-
];
35
+
postInstall = ''
36
+
mkdir $testsout
37
+
cp -R examples $testsout/examples
38
+
'';
41
39
42
40
pythonImportsCheck = [
43
41
"optax"
44
42
];
45
43
46
-
disabledTestPaths = [
47
-
# Requires `flax` which depends on `optax` creating circular dependency.
48
-
"optax/_src/equivalence_test.py"
49
-
# See https://github.com/deepmind/optax/issues/323.
50
-
"examples/lookahead_mnist_test.py"
51
-
];
44
+
# check in passthru.tests.pytest to escape infinite recursion with flax
45
+
doCheck = false;
46
+
47
+
passthru.tests = {
48
+
pytest = callPackage ./tests.nix { };
49
+
};
52
50
53
51
meta = with lib; {
54
52
description = "Optax is a gradient processing and optimization library for JAX.";
+35
pkgs/development/python-modules/optax/tests.nix
+35
pkgs/development/python-modules/optax/tests.nix
···
1
+
{ stdenv
2
+
, buildPythonPackage
3
+
, dm-haiku
4
+
, pytest-xdist
5
+
, pytestCheckHook
6
+
, tensorflow
7
+
, tensorflow-datasets
8
+
, flax
9
+
, optax
10
+
}:
11
+
12
+
buildPythonPackage rec {
13
+
pname = "optax-tests";
14
+
inherit (optax) version;
15
+
16
+
src = optax.testsout;
17
+
18
+
dontBuild = true;
19
+
dontInstall = true;
20
+
21
+
checkInputs = [
22
+
dm-haiku
23
+
pytest-xdist
24
+
pytestCheckHook
25
+
tensorflow
26
+
tensorflow-datasets
27
+
flax
28
+
];
29
+
30
+
disabledTestPaths = [
31
+
# See https://github.com/deepmind/optax/issues/323
32
+
"examples/lookahead_mnist_test.py"
33
+
];
34
+
35
+
}