Add two ensemble sampling methods

docs/source/mcmc.rst

@@ -9,7 +9,9 @@ We provide a high-level overview of the MCMC algorithms in NumPyro:
 * `BarkerMH <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.barker.BarkerMH>`_ is a gradient-based MCMC method that may be competitive with HMC and NUTS for some models. It is applicable to models with continuous latent variables.
 * `HMCGibbs <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.hmc_gibbs.HMCGibbs>`_ combines HMC/NUTS steps with custom Gibbs updates. Gibbs updates must be specified by the user.
 * `DiscreteHMCGibbs <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.hmc_gibbs.DiscreteHMCGibbs>`_ combines HMC/NUTS steps with Gibbs updates for discrete latent variables. The corresponding Gibbs updates are computed automatically.
-* `SA <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.sa.SA>`_ is the only MCMC method in NumPyro that does not leverage gradients. It is only applicable to models with continuous latent variables. It is expected to perform best for models whose latent dimension is low to moderate. It may be a good choice for models with non-differentiable log densities. Note that SA generally requires a *very* large number of samples, as mixing tends to be slow. On the plus side individual steps can be fast.
+* `SA <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.sa.SA>`_ does not leverage gradients. It is only applicable to models with continuous latent variables. It is expected to perform best for models whose latent dimension is low to moderate. It may be a good choice for models with non-differentiable log densities. Note that SA generally requires a *very* large number of samples, as mixing tends to be slow. On the plus side individual steps can be fast.
+* `AIES <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.ensemble.AIES>`_ is a gradient-free ensemble method that informs Metropolis-Hastings proposals by sharing information between chains. It is only applicable to models with continuous latent variables. It is expected to perform best for models whose latent dimension is low to moderate. It may be a good choice for models with non-differentiable log densities, and can be robust to likelihood-free models. AIES generally requires the number of chains to be twice as large as the number of latent parameters, (and ideally larger). 
+* `ESS <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.ensemble.ESS>`_ is a gradient-free ensemble method that shares information between chains to find good slice sampling directions. It tends to be more sample efficient than AIES. It is only applicable to models with continuous latent variables. It is expected to perform best for models whose latent dimension is low to moderate and may be a good choice for models with non-differentiable log densities. ESS generally requires the number of chains to be twice as large as the number of latent parameters, (and ideally larger). 
 
 Like HMC/NUTS, all remaining MCMC algorithms support enumeration over discrete latent variables if possible (see `restrictions <https://pyro.ai/examples/enumeration.html#Restriction-1:-conditional-independence>`_). Enumerated sites need to be marked with `infer={'enumerate': 'parallel'}` like in the `annotation example <https://num.pyro.ai/en/stable/examples/annotation.html>`_.
 
@@ -101,6 +103,30 @@ SA
     :show-inheritance:
     :member-order: bysource
 
+EnsembleSampler
+^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.infer.ensemble.EnsembleSampler
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+AIES
+^^^^
+.. autoclass:: numpyro.infer.ensemble.AIES
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+ESS
+^^^
+.. autoclass:: numpyro.infer.ensemble.ESS
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 .. autofunction:: numpyro.infer.hmc.hmc
 
 .. autofunction:: numpyro.infer.hmc.hmc.init_kernel
@@ -117,6 +143,12 @@ SA
 
 .. autodata:: numpyro.infer.sa.SAState
 
+.. autodata:: numpyro.infer.ensemble.EnsembleSamplerState
+
+.. autodata:: numpyro.infer.ensemble.AIESState
+
+.. autodata:: numpyro.infer.ensemble.ESSState
+
 
 TensorFlow Kernels
 ------------------

numpyro/infer/__init__.py

@@ -10,6 +10,7 @@
     TraceGraph_ELBO,
     TraceMeanField_ELBO,
 )
+from numpyro.infer.ensemble import AIES, ESS
 from numpyro.infer.hmc import HMC, NUTS
 from numpyro.infer.hmc_gibbs import HMCECS, DiscreteHMCGibbs, HMCGibbs
 from numpyro.infer.initialization import (
@@ -29,6 +30,7 @@
 from . import autoguide, reparam
 
 __all__ = [
+    "AIES",
     "autoguide",
     "init_to_feasible",
     "init_to_mean",
@@ -41,6 +43,7 @@
     "BarkerMH",
     "DiscreteHMCGibbs",
     "ELBO",
+    "ESS",
     "HMC",
     "HMCECS",
     "HMCGibbs",