tjh.dev
1{ stdenv
2, buildPythonPackage
3, dm-haiku
4, chex
5, cloudpickle
6, dill
7, dm-tree
8, jaxlib
9, pytest-xdist
10, pytestCheckHook
11, tensorflow
12, bsuite
13, frozendict
14, dm-env
15, scikit-image
16, rlax
17, distrax
18, tensorflow-probability
19, optax }:
20
21buildPythonPackage rec {
22 pname = "dm-haiku-tests";
23 inherit (dm-haiku) version;
24
25 src = dm-haiku.testsout;
26
27 dontBuild = true;
28 dontInstall = true;
29
30 nativeCheckInputs = [
31 bsuite
32 chex
33 cloudpickle
34 dill
35 distrax
36 dm-env
37 dm-haiku
38 dm-tree
39 frozendict
40 jaxlib
41 pytest-xdist
42 pytestCheckHook
43 optax
44 rlax
45 scikit-image
46 tensorflow
47 tensorflow-probability
48 ];
49
50 disabledTests = [
51 # See https://github.com/deepmind/dm-haiku/issues/366.
52 "test_jit_Recurrent"
53 # Assertion errors
54 "test_connect_conv_padding_function_same0"
55 "test_connect_conv_padding_function_valid0"
56 "test_connect_conv_padding_function_same1"
57 "test_connect_conv_padding_function_same2"
58 "test_connect_conv_padding_function_valid1"
59 "test_connect_conv_padding_function_valid2"
60 "test_invalid_axis_ListString"
61 "test_invalid_axis_String"
62 "test_simple_case"
63 "test_simple_case_with_scale"
64 "test_slice_axis"
65 "test_zero_inputs"
66 ];
67
68}