Merge pull request #246761 from GaetanLepage/numpyro

python310Packages.numpyro: 0.11.0 -> 0.12.1

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

@@ -0,0 +1,78 @@
+{ lib
+, buildPythonPackage
+, pythonOlder
+, fetchFromGitHub
+, makefun
+, multipledispatch
+, numpy
+, opt-einsum
+, typing-extensions
+, pyro-ppl
+, torch
+, pandas
+, pillow
+, pyro-api
+, pytest
+, pytest-xdist
+, requests
+, scipy
+, torchvision
+}:
+
+buildPythonPackage rec {
+  pname = "funsor";
+  version = "0.4.5";
+
+  disabled = pythonOlder "3.7";
+
+  src = fetchFromGitHub {
+    owner = "pyro-ppl";
+    repo = "funsor";
+    rev = "refs/tags/${version}";
+    hash = "sha256-3JxnRlJN/kYstcIKt8cvDmqnp8oyoA4jkMBDVuAMa5Q=";
+  };
+
+  # Disable the tests that rely on downloading assets from the internet as well as the linting checks.
+  patches = [
+    ./patch-makefile-for-tests.patch
+  ];
+
+  propagatedBuildInputs = [
+    makefun
+    multipledispatch
+    numpy
+    opt-einsum
+    typing-extensions
+  ];
+
+  nativeCheckInputs = [
+    # Backend
+    pyro-ppl
+    torch
+
+    pandas
+    pillow
+    pyro-api
+    pytest
+    pytest-xdist
+    requests
+    scipy
+    torchvision
+  ];
+
+  # Use the included Makefile to run the tests.
+  checkPhase = ''
+    export FUNSOR_BACKEND=torch
+    make test
+  '';
+
+  pythonImportsCheck = [ "funsor" ];
+
+  meta = with lib; {
+    description = "Functional tensors for probabilistic programming";
+    homepage = "https://funsor.pyro.ai";
+    changelog = "https://github.com/pyro-ppl/funsor/releases/tag/${version}";
+    license = licenses.asl20;
+    maintainers = with maintainers; [ GaetanLepage ];
+  };
+}

pkgs/development/python-modules/funsor/patch-makefile-for-tests.patch

@@ -0,0 +1,27 @@
+diff --git a/Makefile b/Makefile
+index d4b862c..755dfe3 100644
+--- a/Makefile
++++ b/Makefile
+@@ -23,7 +23,7 @@ format: license FORCE
+ 	black .
+ 	isort .
+ 
+-test: lint FORCE
++test: FORCE
+ ifeq (${FUNSOR_BACKEND}, torch)
+ 	pytest -v -n auto test/
+ 	FUNSOR_DEBUG=1 pytest -v test/test_gaussian.py
+@@ -45,10 +45,10 @@ ifeq (${FUNSOR_BACKEND}, torch)
+ 	python examples/slds.py -n 2 -t 50
+ 	python examples/pcfg.py --size 3
+ 	python examples/talbot.py -n 2
+-	python examples/vae.py --smoke-test
++	# python examples/vae.py --smoke-test
+ 	python examples/eeg_slds.py --num-steps 2 --fon --test
+-	python examples/mixed_hmm/experiment.py -d seal -i discrete -g discrete -zi --smoke
+-	python examples/mixed_hmm/experiment.py -d seal -i discrete -g discrete -zi --parallel --smoke
++	# python examples/mixed_hmm/experiment.py -d seal -i discrete -g discrete -zi --smoke
++	# python examples/mixed_hmm/experiment.py -d seal -i discrete -g discrete -zi --parallel --smoke
+ 	python examples/sensor.py --seed=0 --num-frames=2 -n 1
+ 	python examples/adam.py --num-steps=21
+ 	@echo PASS

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

@@ -1,39 +1,41 @@
 { lib
 , buildPythonPackage
+, pythonOlder
 , fetchPypi
 , jax
 , jaxlib
 , multipledispatch
 , numpy
-, pytestCheckHook
-, pythonOlder
-, tensorflow-probability
 , tqdm
+, funsor
+, pytestCheckHook
+, tensorflow-probability
 }:
 
 buildPythonPackage rec {
   pname = "numpyro";
-  version = "0.11.0";
+  version = "0.12.1";
   format = "setuptools";
 
-  disabled = pythonOlder "3.7";
+  disabled = pythonOlder "3.8";
 
   src = fetchPypi {
     inherit version pname;
-    hash = "sha256-01fdGgFZ+G1FwjNwitM6PT1TQx0FtLvs4dBorkFoqo4=";
+    hash = "sha256-S3ifL/KPOJQcyBEYoE1XGxPLmSfh1uT9wJG/YtABBKQ=";
   };
 
   propagatedBuildInputs = [
     jax
     jaxlib
-    numpy
     multipledispatch
+    numpy
     tqdm
   ];
 
   nativeCheckInputs = [
-    tensorflow-probability
+    funsor
     pytestCheckHook
+    tensorflow-probability
   ];
 
   pythonImportsCheck = [

pkgs/top-level/python-packages.nix

@@ -4046,6 +4046,8 @@ self: super: with self; {
 
   funcy = callPackage ../development/python-modules/funcy { };
 
+  funsor = callPackage ../development/python-modules/funsor { };
+
   furl = callPackage ../development/python-modules/furl { };
 
   furo = callPackage ../development/python-modules/furo { };