Merge pull request #59363 from timokau/graph_nets-init

python.pkgs.{tensorflow-probability,dm-sonnet,graph_nets}: init

pkgs/development/python-modules/dm-sonnet/default.nix

@@ -0,0 +1,85 @@
+{ lib
+, fetchFromGitHub
+, buildBazelPackage
+, buildPythonPackage
+, git
+, python
+, six
+, absl-py
+, semantic-version
+, contextlib2
+, wrapt
+, tensorflow
+, tensorflow-probability
+}:
+
+let
+  version = "1.30";
+
+  # first build all binaries and generate setup.py using bazel
+  bazel-build = buildBazelPackage rec {
+    name = "dm-sonnet-bazel-${version}";
+
+    src = fetchFromGitHub {
+      owner = "deepmind";
+      repo = "sonnet";
+      rev = "v${version}";
+      sha256 = "1dli4a4arx2gmb4p676pfibvnpag9f13znisrk9381g7xpqqmaw6";
+    };
+
+    nativeBuildInputs = [
+      git # needed to fetch the bazel deps (protobuf in particular)
+    ];
+
+    # see https://github.com/deepmind/sonnet/blob/master/docs/INSTALL.md
+    bazelTarget = ":install";
+
+    fetchAttrs = {
+      sha256 = "1qwq6xp6gdmy8f0k96q3nsgqs56adrbgq39g5smwik6griwfx9mr";
+    };
+
+    buildAttrs = {
+      preBuild = ''
+        patchShebangs .
+      '';
+
+      installPhase = ''
+        # do not generate a wheel, instead just copy the generated files to $out to be installed by buildPythonPackage
+        sed -i 's,.*bdist_wheel.*,cp -rL . "$out"; exit 0,' bazel-bin/install
+
+        # the target directory "dist" does not actually matter since we're not generating a wheel
+        bazel-bin/install dist
+      '';
+    };
+  };
+
+# now use pip to install the package prepared by bazel
+in buildPythonPackage rec {
+  pname = "dm-sonnet";
+  inherit version;
+
+  src = bazel-build;
+
+  propagatedBuildInputs = [
+    six
+    absl-py
+    semantic-version
+    contextlib2
+    wrapt
+    tensorflow
+    tensorflow-probability
+  ];
+
+  # not sure how to properly run the real test suite -- through bazel?
+  checkPhase = ''
+    ${python.interpreter} -c "import sonnet"
+  '';
+
+  meta = with lib; {
+    description = "TensorFlow-based neural network library";
+    homepage = https://sonnet.dev;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ timokau ];
+    platforms = platforms.linux;
+  };
+}

pkgs/development/python-modules/graph_nets/default.nix

@@ -0,0 +1,47 @@
+{ lib
+, buildPythonPackage
+, fetchPypi
+, tensorflow
+, absl-py
+, dm-sonnet
+, networkx
+, numpy
+, setuptools
+, six
+, future
+}:
+
+buildPythonPackage rec {
+  pname = "graph_nets";
+  version = "1.0.3";
+
+  src = fetchPypi {
+    inherit pname version;
+    sha256 = "15cbs9smmgqz2n9mnlzdbqj3iv9iw179d2g0f9lnimdy7xl4jqdf";
+  };
+
+  buildInputs = [];
+
+  postPatch = ''
+    # https://github.com/deepmind/graph_nets/issues/63
+    sed -i 's/dm-sonnet==1.23/dm-sonnet/' setup.py
+  '';
+
+  propagatedBuildInputs = [
+    tensorflow
+    absl-py
+    dm-sonnet
+    networkx
+    numpy
+    setuptools
+    six
+    future
+  ];
+
+  meta = with lib; {
+    description = "Build Graph Nets in Tensorflow";
+    homepage = https://github.com/deepmind/graph_nets;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ timokau ];
+  };
+}

pkgs/development/python-modules/tensorflow-probability/default.nix

@@ -0,0 +1,41 @@
+{ lib
+, fetchFromGitHub
+, buildPythonPackage
+, tensorflow
+, pytest
+}:
+
+buildPythonPackage rec {
+  pname = "tensorflow-probability";
+  version = "0.6.0";
+
+  src = fetchFromGitHub {
+    owner = "tensorflow";
+    repo = "probability";
+    rev = "v${version}";
+    sha256 = "1y210n4asv8j39pk68bdfrz01gddflvzhxbcvj5jw6rjgaagnhvx";
+  };
+
+  propagatedBuildInputs = [
+    tensorflow
+  ];
+
+  checkInputs = [
+    pytest
+  ];
+
+  # Tests have an invalid import (`tensorflow_probability.opensource`), should
+  # be resolved in the next version with
+  # https://github.com/tensorflow/probability/commit/77d5957f2f0bdddcb46582799cd9c5c5167a1a40
+  doCheck = false;
+  checkPhase = ''
+    py.test
+  '';
+
+  meta = with lib; {
+    description = "Library for probabilistic reasoning and statistical analysis";
+    homepage = https://www.tensorflow.org/probability/;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ timokau ];
+  };
+}

pkgs/top-level/python-packages.nix

@@ -4877,6 +4877,8 @@ in {
 
   graphite_beacon = callPackage ../development/python-modules/graphite_beacon { };
 
+  graph_nets = callPackage ../development/python-modules/graph_nets { };
+
   influxgraph = callPackage ../development/python-modules/influxgraph { };
 
   graphitepager = callPackage ../development/python-modules/graphitepager { };
@@ -4927,6 +4929,8 @@ in {
 
   snapperGUI = callPackage ../development/python-modules/snappergui { };
 
+  dm-sonnet = callPackage ../development/python-modules/dm-sonnet { };
+
   uncertainties = callPackage ../development/python-modules/uncertainties { };
 
   funcy = callPackage ../development/python-modules/funcy { };
@@ -5143,6 +5147,8 @@ in {
 
   tensorflow-estimator = callPackage ../development/python-modules/tensorflow-estimator { };
 
+  tensorflow-probability = callPackage ../development/python-modules/tensorflow-probability { };
+
   tensorflow-tensorboard = callPackage ../development/python-modules/tensorflow-tensorboard { };
 
   tensorflow =