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