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