Merge branch 'master' into adjoint-leak

pyro-ppl · Mar 10, 2025 · c9fd4c5 · c9fd4c5
2 parents f0fe41a + 7365f38
commit c9fd4c5
Show file tree

Hide file tree

Showing 25 changed files with 1,974 additions and 117 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -39,9 +39,11 @@ jobs:
       run: |
         make lint
     - name: Build documentation
+      if: matrix.python-version != '3.9'
       run: |
         make docs
     - name: Test documentation
+      if: matrix.python-version != '3.9'
       run: |
         make doctest
         python -m doctest -v README.md
@@ -83,12 +85,10 @@ jobs:
       env:
         JAX_CHECK_TRACER_LEAKS: 1
       run: |
-        pytest -vs test/contrib/einstein/test_steinvi.py::test_run_smoke -k ASVGD
         pytest -vs test/infer/test_mcmc.py::test_chain_inside_jit
         pytest -vs test/infer/test_mcmc.py::test_chain_jit_args_smoke
         pytest -vs test/infer/test_mcmc.py::test_reuse_mcmc_run
         pytest -vs test/infer/test_mcmc.py::test_model_with_multiple_exec_paths
-        pytest -vs test/infer/test_svi.py::test_mutable_state
         pytest -vs test/test_distributions.py::test_mean_var -k Gompertz
 
     - name: Coveralls

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -35,4 +35,4 @@ repos:
       - id: codespell
         stages: [pre-commit, commit-msg]
         args:
-          [--ignore-words-list, "Teh,aas,ans", --check-filenames, --skip, "*.ipynb"]
+          [--ignore-words-list, "Teh,aas,ans,dout", --check-filenames, --skip, "*.ipynb"]
diff --git a/docs/source/_static/img/tutorials/nnx_example.png b/docs/source/_static/img/tutorials/nnx_example.png
diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
@@ -136,6 +136,14 @@ Chi2
     :show-inheritance:
     :member-order: bysource
 
+CirculantNormal
+^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.continuous.CirculantNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 Dirichlet
 ^^^^^^^^^
 .. autoclass:: numpyro.distributions.continuous.Dirichlet
@@ -998,6 +1006,14 @@ OrderedTransform
     :show-inheritance:
     :member-order: bysource
 
+PackRealFastFourierCoefficientsTransform
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.transforms.PackRealFastFourierCoefficientsTransform
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 PermuteTransform
 ^^^^^^^^^^^^^^^^
 .. autoclass:: numpyro.distributions.transforms.PermuteTransform

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -39,6 +39,8 @@ NumPyro documentation
    tutorials/censoring
    tutorials/hsgp_example
    tutorials/other_samplers
+   tutorials/circulant_gp
+   tutorials/nnx_example
 
 .. nbgallery::
    :maxdepth: 1

diff --git a/docs/source/primitives.rst b/docs/source/primitives.rst
@@ -51,6 +51,10 @@ haiku_module
 ------------
 .. autofunction:: numpyro.contrib.module.haiku_module
 
+nnx_module
+----------
+.. autofunction:: numpyro.contrib.module.nnx_module
+
 random_flax_module
 ------------------
 .. autofunction:: numpyro.contrib.module.random_flax_module
@@ -59,6 +63,10 @@ random_haiku_module
 -------------------
 .. autofunction:: numpyro.contrib.module.random_haiku_module
 
+random_nnx_module
+-----------------
+.. autofunction:: numpyro.contrib.module.random_nnx_module
+
 scan
 ----
 .. autofunction:: numpyro.contrib.control_flow.scan

diff --git a/notebooks/source/circulant_gp.ipynb b/notebooks/source/circulant_gp.ipynb
diff --git a/notebooks/source/nnx_example.ipynb b/notebooks/source/nnx_example.ipynb