pyrox.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 pythonOlder,
22}:
23
24buildPythonPackage rec {
25 pname = "chex";
26 version = "0.1.89";
27 pyproject = true;
28
29 src = fetchFromGitHub {
30 owner = "deepmind";
31 repo = "chex";
32 tag = "v${version}";
33 hash = "sha256-eTEfmbpmwtCuphtOq0iHUT8zGfyQ4/aUorL4FQwcKBw=";
34 };
35
36 build-system = [ setuptools ];
37
38 dependencies = [
39 absl-py
40 jax
41 jaxlib
42 numpy
43 toolz
44 typing-extensions
45 ];
46
47 pythonImportsCheck = [ "chex" ];
48
49 nativeCheckInputs = [
50 cloudpickle
51 dm-tree
52 pytestCheckHook
53 ];
54
55 meta = {
56 description = "Library of utilities for helping to write reliable JAX code";
57 homepage = "https://github.com/deepmind/chex";
58 changelog = "https://github.com/google-deepmind/chex/releases/tag/v${version}";
59 license = lib.licenses.asl20;
60 maintainers = with lib.maintainers; [ ndl ];
61 };
62}