1{ buildPythonPackage
2, dm-haiku
3, pytest-xdist
4, pytestCheckHook
5, tensorflow
6, tensorflow-datasets
7, flax
8, optax
9}:
10
11buildPythonPackage {
12 pname = "optax-tests";
13 inherit (optax) version;
14
15 src = optax.testsout;
16
17 dontBuild = true;
18 dontInstall = true;
19
20 nativeCheckInputs = [
21 dm-haiku
22 pytest-xdist
23 pytestCheckHook
24 tensorflow
25 tensorflow-datasets
26 flax
27 ];
28
29 disabledTestPaths = [
30 # See https://github.com/deepmind/optax/issues/323
31 "examples/lookahead_mnist_test.py"
32 ];
33}