pyrox.dev
1{ buildPythonPackage
2, fetchFromGitHub
3, fetchpatch
4, jax
5, jaxlib
6, lib
7, poetry-core
8}:
9
10buildPythonPackage rec {
11 pname = "treeo";
12 # Note that there is a version 0.4.0, but it was released in error. At the
13 # time of writing (2022-03-29), v0.0.11 is the latest as reported on GitHub
14 # and PyPI.
15 version = "0.0.11";
16 format = "pyproject";
17
18 src = fetchFromGitHub {
19 owner = "cgarciae";
20 repo = pname;
21 rev = version;
22 hash = "sha256-zs3F8i+G5OX/A9wOO60xVuvnm2QqrL+dHIrC0qwH37o=";
23 };
24
25 # See https://github.com/cgarciae/treex/issues/68.
26 patches = [
27 (fetchpatch {
28 url = "https://github.com/cgarciae/treeo/pull/14/commits/022915da2b3bf76406a7c79d1b4593bee7956f16.patch";
29 hash = "sha256-WGxJqqrf2g0yZe30RyG1xxbloiqj1awuf1Y4eh5y+z0=";
30 })
31 (fetchpatch {
32 url = "https://github.com/cgarciae/treeo/pull/14/commits/99f9488bd0c977780844fd79743167b0010d359b.patch";
33 hash = "sha256-oKDYs+Ah0QXkhiJysIudQ6VLIiUiIcnQisxYp6GJuTc=";
34 })
35 ];
36
37 nativeBuildInputs = [
38 poetry-core
39 ];
40
41 # jax is not declared in the dependencies, but is necessary.
42 propagatedBuildInputs = [
43 jax
44 ];
45
46 checkInputs = [ jaxlib ];
47 pythonImportsCheck = [
48 "treeo"
49 ];
50
51 meta = with lib; {
52 description = "A small library for creating and manipulating custom JAX Pytree classes.";
53 homepage = "https://github.com/cgarciae/treeo";
54 license = licenses.mit;
55 maintainers = with maintainers; [ ndl ];
56 };
57}