tjh.dev
1{ lib
2, fetchFromGitHub
3, buildBazelPackage
4, buildPythonPackage
5, git
6, python
7, six
8, absl-py
9, semantic-version
10, contextlib2
11, wrapt
12, tensorflow
13, tensorflow-probability
14, tensorflow-estimator
15}:
16
17let
18 version = "1.33";
19
20 # first build all binaries and generate setup.py using bazel
21 bazel-build = buildBazelPackage {
22 name = "dm-sonnet-bazel-${version}";
23
24 src = fetchFromGitHub {
25 owner = "deepmind";
26 repo = "sonnet";
27 rev = "v${version}";
28 sha256 = "1nqsja1s8jrkq6v1whgh7smk17313mjr9vs3k5c1m8px4yblzhqc";
29 };
30
31 nativeBuildInputs = [
32 git # needed to fetch the bazel deps (protobuf in particular)
33 ];
34
35 # see https://github.com/deepmind/sonnet/blob/master/docs/INSTALL.md
36 bazelTarget = ":install";
37
38 fetchAttrs = {
39 sha256 = "0mxma7jajm42v1hv6agl909xra0azihj588032ivhlmmh403x6wg";
40 };
41
42 bazelFlags = [
43 # https://github.com/deepmind/sonnet/issues/134
44 "--incompatible_disable_deprecated_attr_params=false"
45 ];
46
47 buildAttrs = {
48 preBuild = ''
49 patchShebangs .
50 '';
51
52 installPhase = ''
53 # do not generate a wheel, instead just copy the generated files to $out to be installed by buildPythonPackage
54 sed -i 's,.*bdist_wheel.*,cp -rL . "$out"; exit 0,' bazel-bin/install
55
56 # the target directory "dist" does not actually matter since we're not generating a wheel
57 bazel-bin/install dist
58 '';
59 };
60 };
61
62# now use pip to install the package prepared by bazel
63in buildPythonPackage {
64 pname = "dm-sonnet";
65 inherit version;
66
67 src = bazel-build;
68
69 propagatedBuildInputs = [
70 six
71 absl-py
72 semantic-version
73 contextlib2
74 wrapt
75 tensorflow
76 tensorflow-probability
77 tensorflow-estimator
78 ];
79
80 # not sure how to properly run the real test suite -- through bazel?
81 checkPhase = ''
82 ${python.interpreter} -c "import sonnet"
83 '';
84
85 meta = with lib; {
86 description = "TensorFlow-based neural network library";
87 homepage = https://sonnet.dev;
88 license = licenses.asl20;
89 maintainers = with maintainers; [ timokau ];
90 platforms = platforms.linux;
91 };
92}