1{
2 lib,
3 buildPythonPackage,
4 pythonOlder,
5 fetchFromGitHub,
6 setuptools,
7 absl-py,
8 cloudpickle,
9 dm-tree,
10 jax,
11 jaxlib,
12 numpy,
13 pytestCheckHook,
14 toolz,
15 typing-extensions,
16}:
17
18buildPythonPackage rec {
19 pname = "chex";
20 version = "0.1.86";
21 pyproject = true;
22
23 disabled = pythonOlder "3.9";
24
25 src = fetchFromGitHub {
26 owner = "deepmind";
27 repo = "chex";
28 rev = "refs/tags/v${version}";
29 hash = "sha256-Z5Ns4fG5pC99I4xdGjDMKX6YZpTtd1y0TWcIOtr7dug=";
30 };
31
32 build-system = [ setuptools ];
33
34 dependencies = [
35 absl-py
36 jaxlib
37 jax
38 numpy
39 toolz
40 typing-extensions
41 ];
42
43 pythonImportsCheck = [ "chex" ];
44
45 nativeCheckInputs = [
46 cloudpickle
47 dm-tree
48 pytestCheckHook
49 ];
50
51 meta = with lib; {
52 description = "Chex is a library of utilities for helping to write reliable JAX code.";
53 homepage = "https://github.com/deepmind/chex";
54 changelog = "https://github.com/google-deepmind/chex/releases/tag/v${version}";
55 license = licenses.asl20;
56 maintainers = with maintainers; [ ndl ];
57 };
58}