1 files changed, 46 insertions, 15 deletions
diff --git a/pkgs/development/python-modules/pyro-ppl/default.nix b/pkgs/development/python-modules/pyro-ppl/default.nix
index d4837dd9fb1..7bc6ac00925 100644
--- a/pkgs/development/python-modules/pyro-ppl/default.nix
+++ b/pkgs/development/python-modules/pyro-ppl/default.nix
@@ -1,29 +1,62 @@
-{ buildPythonPackage, fetchPypi, lib, pytorch, contextlib2
-, graphviz, networkx, six, opt-einsum, tqdm, pyro-api }:
+{ lib
+, buildPythonPackage
+, fetchPypi
+, graphviz
+, jupyter
+, matplotlib
+, networkx
+, opt-einsum
+, pandas
+, pillow
+, pyro-api
+, pythonOlder
+, pytorch
+, scikit-learn
+, seaborn
+, torchvision
+, tqdm
+, wget
+}:
 
 buildPythonPackage rec {
-  version = "1.8.1";
   pname = "pyro-ppl";
+  version = "1.8.1";
+  format = "setuptools";
+
+  disabled = pythonOlder "3.7";
 
   src = fetchPypi {
     inherit version pname;
-    sha256 = "sha256-18BJ6y50haYStN2ZwkwwnMhgx8vGsZczhwNPVDbRyNY=";
+    hash = "sha256-18BJ6y50haYStN2ZwkwwnMhgx8vGsZczhwNPVDbRyNY=";
   };
 
   propagatedBuildInputs = [
     pyro-api
     pytorch
-    contextlib2
-    # TODO(tom): graphviz pulls in a lot of dependencies - make
-    # optional when some time to figure out how.
-    graphviz
     networkx
-    six
     opt-einsum
     tqdm
   ];
 
+  passthru.optional-dependencies = {
+    extras = [
+      graphviz
+      jupyter
+      # lap
+      matplotlib
+      pandas
+      pillow
+      scikit-learn
+      seaborn
+      torchvision
+      # visdom
+      wget
+    ];
+  };
+
   # pyro not shipping tests do simple smoke test instead
+  doCheck = false;
+
   pythonImportsCheck = [
     "pyro"
     "pyro.distributions"
@@ -31,12 +64,10 @@ buildPythonPackage rec {
     "pyro.optim"
   ];
 
-  doCheck = false;
-
-  meta = {
-    description = "A Python library for probabilistic modeling and inference";
+  meta = with lib; {
+    description = "Library for probabilistic modeling and inference";
     homepage = "http://pyro.ai";
-    license = lib.licenses.mit;
-    maintainers = with lib.maintainers; [ teh georgewhewell ];
+    license = licenses.asl20;
+    maintainers = with maintainers; [ teh georgewhewell ];
   };
 }