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