pyrox.dev
1{
2 lib,
3 stdenv,
4 fetchpatch2,
5
6 # bazel wheel
7 buildBazelPackage,
8 fetchFromGitHub,
9
10 # nativeBuildInputs
11 python,
12 setuptools,
13 wheel,
14 absl-py,
15
16 #bazel_6,
17 bazel,
18 cctools,
19
20 # python package
21 buildPythonPackage,
22
23 # dependencies
24 cloudpickle,
25 decorator,
26 dm-tree,
27 gast,
28 keras,
29 numpy,
30 six,
31 tensorflow,
32
33 # tests
34 hypothesis,
35 matplotlib,
36 mock,
37 mpmath,
38 pandas,
39 pytest,
40 scipy,
41}:
42
43let
44 version = "0.25.0";
45 pname = "tensorflow-probability";
46
47 # first build all binaries and generate setup.py using bazel
48 bazel-wheel = buildBazelPackage {
49 name = "tensorflow_probability-${version}-py2.py3-none-any.whl";
50 src = fetchFromGitHub {
51 owner = "tensorflow";
52 repo = "probability";
53 rev = "refs/tags/v${version}";
54 hash = "sha256-LXQfGFgnM7WYUQjJ2Y3jskdeJ/dEKz+Afg+UOQjv5kc=";
55 };
56
57 patches = [
58 # AttributeError: jax.interpreters.xla.pytype_aval_mappings was deprecated in JAX v0.5.0 and
59 # removed in JAX v0.7.0. jax.core.pytype_aval_mappings can be used as a replacement in most cases.
60 # TODO: remove when updating to the next release
61 (fetchpatch2 {
62 name = "future-proof-reference-to-deprecated-pytype_aval_mappings";
63 url = "https://github.com/tensorflow/probability/commit/135080b6b1ac5724fc1731b0a9ca6f2010b1aea5.patch";
64 hash = "sha256-27yWIw5pI86KcUz0TsYwRFyLDoeiqmxgsRMBXaauzVw=";
65 })
66 ];
67
68 nativeBuildInputs = [
69 absl-py
70 # needed to create the output wheel in installPhase
71 python
72 setuptools
73 tensorflow
74 wheel
75 ];
76
77 #bazel = bazel_6;
78 bazel = bazel;
79
80 bazelTargets = [ ":pip_pkg" ];
81 LIBTOOL = lib.optionalString stdenv.hostPlatform.isDarwin "${cctools}/bin/libtool";
82
83 fetchAttrs = {
84 sha256 = "sha256-TbWcWYidyXuAMgBnO2/k0NKCzc4wThf2uUeC3QxdBJY=";
85 };
86
87 buildAttrs = {
88 preBuild = ''
89 patchShebangs .
90 '';
91
92 installPhase = ''
93 # work around timestamp issues
94 # https://github.com/NixOS/nixpkgs/issues/270#issuecomment-467583872
95 export SOURCE_DATE_EPOCH=315532800
96
97 # First build, then move. Otherwise pip_pkg would create the dir $out
98 # and then put the wheel in that directory. However we want $out to
99 # point directly to the wheel file.
100 ./bazel-bin/pip_pkg . --release
101 mv *.whl "$out"
102 '';
103 };
104 };
105in
106buildPythonPackage {
107 inherit version pname;
108 format = "wheel";
109
110 src = bazel-wheel;
111
112 dependencies = [
113 cloudpickle
114 decorator
115 dm-tree
116 gast
117 keras
118 numpy
119 six
120 tensorflow
121 ];
122
123 # Listed here:
124 # https://github.com/tensorflow/probability/blob/f3777158691787d3658b5e80883fe1a933d48989/testing/dependency_install_lib.sh#L83
125 nativeCheckInputs = [
126 hypothesis
127 matplotlib
128 mock
129 mpmath
130 pandas
131 pytest
132 scipy
133 ];
134
135 # Ideally, we run unit tests with pytest, but in checkPhase, only the Bazel-build wheel is available.
136 # But it seems not guaranteed that running the tests with pytest will even work, see
137 # https://github.com/tensorflow/probability/blob/c2a10877feb2c4c06a4dc58281e69c37a11315b9/CONTRIBUTING.md?plain=1#L69
138 # Ideally, tests would be run using Bazel. For now, lets's do a...
139
140 # sanity check
141 pythonImportsCheck = [ "tensorflow_probability" ];
142
143 meta = {
144 description = "Library for probabilistic reasoning and statistical analysis";
145 homepage = "https://www.tensorflow.org/probability/";
146 changelog = "https://github.com/tensorflow/probability/releases/tag/v${version}";
147 license = lib.licenses.asl20;
148 maintainers = with lib.maintainers; [ GaetanLepage ];
149 # Needs update for Bazel 7.
150 broken = true;
151 };
152}