pyrox.dev
1{ absl-py
2, buildPythonPackage
3, dm-tree
4, fetchFromGitHub
5, jax
6, jaxlib
7, lib
8, numpy
9, pytestCheckHook
10, toolz
11}:
12
13buildPythonPackage rec {
14 pname = "chex";
15 version = "0.1.5";
16 format = "setuptools";
17
18 src = fetchFromGitHub {
19 owner = "deepmind";
20 repo = pname;
21 rev = "refs/tags/v${version}";
22 hash = "sha256-FYB0IhokM74HBY2wOJYE6xJrHxHHWhDSAZpWFs6HFu0=";
23 };
24
25 propagatedBuildInputs = [
26 absl-py
27 dm-tree
28 jax
29 numpy
30 toolz
31 ];
32
33 pythonImportsCheck = [
34 "chex"
35 ];
36
37 checkInputs = [
38 jaxlib
39 pytestCheckHook
40 ];
41
42 disabledTests = [
43 # See https://github.com/deepmind/chex/issues/204.
44 "test_uninspected_checks"
45 ];
46
47 meta = with lib; {
48 description = "Chex is a library of utilities for helping to write reliable JAX code.";
49 homepage = "https://github.com/deepmind/chex";
50 license = licenses.asl20;
51 maintainers = with maintainers; [ ndl ];
52 };
53}