1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 flit-core,
8
9 # dependencies
10 numpy,
11
12 # optional-dependencies
13 ipython,
14 jax,
15 palettable,
16
17 # tests
18 absl-py,
19 jaxlib,
20 omegaconf,
21 pydantic,
22 pytestCheckHook,
23 torch,
24}:
25
26buildPythonPackage rec {
27 pname = "treescope";
28 version = "0.1.10";
29 pyproject = true;
30
31 src = fetchFromGitHub {
32 owner = "google-deepmind";
33 repo = "treescope";
34 tag = "v${version}";
35 hash = "sha256-SfycwuI/B7S/rKkaqxtnJI26q89313pvj/Xsomg6qyA=";
36 };
37
38 build-system = [ flit-core ];
39
40 dependencies = [ numpy ];
41
42 optional-dependencies = {
43 notebook = [
44 ipython
45 jax
46 palettable
47 ];
48 };
49
50 pythonImportsCheck = [ "treescope" ];
51
52 nativeCheckInputs = [
53 absl-py
54 jax
55 jaxlib
56 omegaconf
57 pydantic
58 pytestCheckHook
59 torch
60 ];
61
62 meta = {
63 description = "Interactive HTML pretty-printer for machine learning research in IPython notebooks";
64 homepage = "https://github.com/google-deepmind/treescope";
65 changelog = "https://github.com/google-deepmind/treescope/releases/tag/v${version}";
66 license = lib.licenses.asl20;
67 maintainers = with lib.maintainers; [ GaetanLepage ];
68 };
69}