Fix links in examples and tutorials

README.md

@@ -303,7 +303,7 @@ conda install -c conda-forge numpyro
    - Any `torch` operation in your model will need to be written in terms of the corresponding `jax.numpy` operation. Additionally, not all `torch` operations have a `numpy` counterpart (and vice-versa), and sometimes there are minor differences in the API.
    - `pyro.sample` statements outside an inference context will need to be wrapped in a `seed` handler, as mentioned above.
    - There is no global parameter store, and as such using `numpyro.param` outside an inference context will have no effect. To retrieve the optimized parameter values from SVI, use the [SVI.get_params](https://num.pyro.ai/en/latest/svi.html#numpyro.infer.svi.SVI.get_params) method. Note that you can still use `param` statements inside a model and NumPyro will use the [substitute](https://num.pyro.ai/en/latest/handlers.html#substitute) effect handler internally to substitute values from the optimizer when running the model in SVI.
-   - PyTorch neural network modules will need to rewritten as [stax](https://github.com/google/jax#neural-net-building-with-stax) neural networks. See the [VAE](#examples) example for differences in syntax between the two backends.
+   - PyTorch neural network modules will need to rewritten as [stax](https://github.com/google/jax#neural-net-building-with-stax), [flax](https://flax.readthedocs.io/en/latest/), or [haiku](https://dm-haiku.readthedocs.io/en/latest/) neural networks. See the [VAE](https://num.pyro.ai/en/latest/examples/vae.html) and [ProdLDA](https://num.pyro.ai/en/stable/examples/prodlda.html) examples for differences in syntax between the two backends.
    - JAX works best with functional code, particularly if we would like to leverage JIT compilation, which NumPyro does internally for many inference subroutines. As such, if your model has side-effects that are not visible to the JAX tracer, it may need to rewritten in a more functional style.
 
    For most small models, changes required to run inference in NumPyro should be minor. Additionally, we are working on [pyro-api](https://github.com/pyro-ppl/pyro-api) which allows you to write the same code and dispatch it to multiple backends, including NumPyro. This will necessarily be more restrictive, but has the advantage of being backend agnostic. See the [documentation](https://pyro-api.readthedocs.io/en/latest/dispatch.html#module-pyroapi.dispatch) for an example, and let us know your feedback.

docs/requirements.txt

@@ -1,21 +1,22 @@
 dm-haiku
 flax
 funsor
-ipython<=8.6.0  # strict the version for https://github.com/ipython/ipython/issues/13845
+ipython
 jax
 jaxlib
 jaxns==2.2.6
-Jinja2<3.1
+Jinja2
 matplotlib
 multipledispatch
-nbsphinx==0.8.9
+nbsphinx>=0.8.9
 numpy
 optax
 pillow
 pylab-sdk
 pyyaml
+readthedocs-sphinx-search>=0.3.2
 sphinx>=5
-sphinx-gallery==0.10.1  # tutorial page is not rendered correctly with newer versions
-sphinx_rtd_theme==1.0.0
+sphinx-gallery
+sphinx_rtd_theme
 tensorflow_probability
 tqdm

notebooks/source/bayesian_hierarchical_stacking.ipynb

@@ -24,16 +24,17 @@
     "## Table of Contents\n",
     "\n",
     "* [Intro](#intro)\n",
-    "* [1. Exploratory Data Analysis](#1)\n",
-    "* [2. Prepare 6 Different Models](#2)\n",
-    "    * [2.1 Feature Engineering](#2.1)\n",
-    "    * [2.2 Training](#2.2)\n",
-    "* [3. Bayesian Hierarchical Stacking](#3)\n",
-    "    * [3.1 Prepare stacking datasets](#3.1)\n",
-    "    * [3.2 Define stacking model](#3.2)\n",
-    "* [4. Evaluate on test set](#4)\n",
-    "    * [4.1 Stack predictions](#4.1)\n",
-    "    * [4.2 Compare methods](#4.2)\n",
+    "* [1. Exploratory Data Analysis](#1.-Exploratory-Data-Analysis)\n",
+    "* [2. Prepare 6 Different Models](#2.-Prepare-6-different-candidate-models)\n",
+    "    * [2.1 Feature Engineering](#2.1-Feature-Engineering)\n",
+    "    * [2.2 Training](#2.2-Training)\n",
+    "    * [2.3 Estimate leave-one-out cross-validated score for each training point](#2.3-Estimate-leave-one-out-cross-validated-score-for-each-training-point)\n",
+    "* [3. Bayesian Hierarchical Stacking](#3.-Bayesian-Hierarchical-Stacking)\n",
+    "    * [3.1 Prepare stacking datasets](#3.1-Prepare-stacking-datasets)\n",
+    "    * [3.2 Define stacking model](#3.2-Define-stacking-model)\n",
+    "* [4. Evaluate on test set](#4.-Evaluate-on-test-set)\n",
+    "    * [4.1 Stack predictions](#4.1-Stack-predictions)\n",
+    "    * [4.2 Compare methods](#4.2-Compare-methods)\n",
     "* [Conclusion](#conclusion)\n",
     "* [References](#references)\n",
     "\n",

notebooks/source/bayesian_regression.ipynb

@@ -46,9 +46,10 @@
     "2. [Regression Model to Predict Divorce Rate](#Regression-Model-to-Predict-Divorce-Rate)\n",
     "    - [Model-1: Predictor-Marriage Rate](#Model-1:-Predictor---Marriage-Rate)\n",
     "       - [Posterior Distribution over the Regression Parameters](#Posterior-Distribution-over-the-Regression-Parameters)\n",
+    "       - [Prior Predictive Distribution](#Prior-Predictive-Distribution)\n",
     "       - [Posterior Predictive Distribution](#Posterior-Predictive-Distribution)\n",
     "       - [Predictive Utility With Effect Handlers](#Predictive-Utility-With-Effect-Handlers)\n",
-    "       - [Model Predictive Density](#Model-Predictive-Density)\n",
+    "       - [Posterior Predictive Density](#Posterior-Predictive-Density)\n",
     "    - [Model-2: Predictor-Median Age of Marriage](#Model-2:-Predictor---Median-Age-of-Marriage)\n",
     "    - [Model-3: Predictor-Marriage Rate and Median Age of Marriage](#Model-3:-Predictor---Marriage-Rate-and-Median-Age-of-Marriage)\n",
     "    - [Divorce Rate Residuals by State](#Divorce-Rate-Residuals-by-State)\n",

notebooks/source/truncated_distributions.ipynb

@@ -15,16 +15,16 @@
     "\n",
     "### Table of contents\n",
     "\n",
-    "* [0. Setup](#0)\n",
-    "* [1. What is a truncated distribution?](#1)\n",
-    "* [2. What is a folded distribution?](#2)\n",
-    "* [3. Sampling from truncated and folded distributions](#3)\n",
-    "* [4. Ready-to-use truncated and folded distributions](#4)\n",
-    "* [5. Building your own truncanted distributions](#5)\n",
-    "    * [5.1 Recap of NumPyro distributions](#5.1)\n",
-    "    * [5.2 Right-truncated normal](#5.2)\n",
-    "    * [5.3 Left-truncated Poisson](#5.3)\n",
-    "* [6. References and related material](#references)\n",
+    "* [0. Setup](#Setup)\n",
+    "* [1. What is a truncated distribution?](#1.-What-are-truncated-distributions?)\n",
+    "* [2. What is a folded distribution?](#2.-What-is-a-folded-distribution?)\n",
+    "* [3. Sampling from truncated and folded distributions](#3.-Sampling-from-truncated-and-folded-distributions)\n",
+    "* [4. Ready-to-use truncated and folded distributions](#4.-Ready-to-use-truncated-and-folded-distributions)\n",
+    "* [5. Building your own truncanted distributions](#5.-Building-your-own-truncated-distribution)\n",
+    "    * [5.1 Recap of NumPyro distributions](#5.1-Recap-of-NumPyro-distributions)\n",
+    "    * [5.2 Right-truncated normal](#5.2-Example:-Right-truncated-normal)\n",
+    "    * [5.3 Left-truncated Poisson](#5.3-Example:-Left-truncated-Poisson)\n",
+    "* [6. References and related material](#References-and-related-material)\n",
     "\n",
     "\n",
     "### Setup  <a class=\"anchor\" id=\"0\"></a>\n",
@@ -164,7 +164,7 @@
     "    <img src=\"https://i.ibb.co/3d2xJbc/folded-normal.png\" alt=\"folded\" width=\"900\"/>\n",
     "</figure>\n",
     "\n",
-    "As you can see, the resulting distribution is different from the truncated case. In particular, the density ratio between points, $p(a)/p(b)$, is in general not the same after folding. For some examples in which folding is relevant see [references 3 and 4](#references)\n",
+    "As you can see, the resulting distribution is different from the truncated case. In particular, the density ratio between points, $p(a)/p(b)$, is in general not the same after folding. For some examples in which folding is relevant see [references 3 and 4](#References-and-related-material)\n",
     "\n",
     "If the original distribution is symmetric around zero, then folding and truncating at zero have the same effect."
    ]

setup.py

@@ -42,14 +42,13 @@
     extras_require={
         "doc": [
             "ipython",  # sphinx needs this to render codes
-            "nbsphinx>=0.8.5",
-            "readthedocs-sphinx-search==0.1.0",
+            "nbsphinx>=0.8.9",
+            "readthedocs-sphinx-search>=0.3.2",
             "sphinx>=5",
             "sphinx_rtd_theme",
             "sphinx-gallery",
         ],
         "test": [
-            "importlib-metadata<5.0",
             "importlib-metadata<5.0",
             "ruff>=0.1.8",
             "pytest>=4.1",