tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{ absl-py
2, buildPythonPackage
3, dm-tree
4, fetchFromGitHub
5, jax
6, jaxlib
7, lib
8, numpy
9, pytestCheckHook
10, toolz
11}:
12
13buildPythonPackage rec {
14 pname = "chex";
15 version = "0.1.2";
16 format = "setuptools";
17
18 src = fetchFromGitHub {
19 owner = "deepmind";
20 repo = pname;
21 rev = "v${version}";
22 hash = "sha256-NtZYOHByKBcKmhRaNULwaQqxfoPRmgbtJ3cFHNfy4E8=";
23 };
24
25 propagatedBuildInputs = [
26 absl-py
27 dm-tree
28 jax
29 numpy
30 toolz
31 ];
32
33 pythonImportsCheck = [
34 "chex"
35 ];
36
37 checkInputs = [
38 jaxlib
39 pytestCheckHook
40 ];
41
42 meta = with lib; {
43 description = "Chex is a library of utilities for helping to write reliable JAX code.";
44 homepage = "https://github.com/deepmind/chex";
45 license = licenses.asl20;
46 maintainers = with maintainers; [ ndl ];
47 };
48}