Make Gibbs kernels composable

.github/workflows/ci.yml

@@ -146,7 +146,7 @@ jobs:
         run: |
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/contrib/stochastic_support/test_dcc.py
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/contrib/test_tfp.py -k "chain"
-          XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_hmc_gibbs.py -k "chain"
+          XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_hmc_gibbs.py test/infer/test_gibbs.py -k "chain"
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/infer/test_mcmc.py -k "chain or pmap or vmap"
           XLA_FLAGS="--xla_force_host_platform_device_count=2" uv run pytest -vs test/test_compile.py -k "chain"
       - name: Test custom prng

docs/source/mcmc.rst

@@ -9,6 +9,7 @@ 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.
+* `Gibbs <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.gibbs.Gibbs>`_ composes any number of block kernels (HMC/NUTS, `DiscreteGibbs <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.gibbs.DiscreteGibbs>`_, `CustomGibbs <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.gibbs.CustomGibbs>`_, or nested `Gibbs`), each owning a subset of the latent variables and conditioned on the others. `HMCGibbs` and `DiscreteHMCGibbs` are two-block instances of it.
 * `SA <https://num.pyro.ai/en/latest/mcmc.html#numpyro.infer.sa.SA>`_ is a gradient-free MCMC method. 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 MCMC 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 MCMC 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). 
@@ -57,6 +58,30 @@ NUTS
     :show-inheritance:
     :member-order: bysource
 
+Gibbs
+^^^^^
+.. autoclass:: numpyro.infer.gibbs.Gibbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+CustomGibbs
+^^^^^^^^^^^
+.. autoclass:: numpyro.infer.gibbs.CustomGibbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+DiscreteGibbs
+^^^^^^^^^^^^^
+.. autoclass:: numpyro.infer.gibbs.DiscreteGibbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 HMCGibbs
 ^^^^^^^^
 .. autoclass:: numpyro.infer.hmc_gibbs.HMCGibbs
@@ -133,7 +158,13 @@ ESS
 
 .. autodata:: numpyro.infer.hmc.HMCState
 
-.. autodata:: numpyro.infer.hmc_gibbs.HMCGibbsState
+.. autoclass:: numpyro.infer.gibbs.GibbsState
+
+.. autoclass:: numpyro.infer.gibbs.CustomGibbsState
+
+.. autoclass:: numpyro.infer.gibbs.DiscreteGibbsState
+
+.. autoclass:: numpyro.infer.hmc_gibbs.HMCGibbsState
 
 .. autodata:: numpyro.infer.sa.SAState
 

numpyro/_typing.py

@@ -36,3 +36,22 @@
 
 
 NumLikeT = TypeVar("NumLikeT", bound=NumLike)
+
+
+ModelArgs: TypeAlias = tuple[Any, ...]
+"""Positional arguments of a model, as passed to ``MCMC.run(rng_key, *args)``."""
+
+ModelKwargs: TypeAlias = dict[str, Any]
+"""Keyword arguments of a model; may carry reserved keys such as ``GIBBS_SITES_KWARG``."""
+
+SiteValues: TypeAlias = dict[str, jax.Array]
+"""Values keyed by site name (a sample, a set of init params, a conditioning set)."""
+
+PotentialFn: TypeAlias = Callable[[SiteValues], jax.Array]
+"""Negative log joint as a function of (unconstrained) site values."""
+
+ConstrainFn: TypeAlias = Callable[[SiteValues], SiteValues]
+"""Maps site values to site values (constrain / postprocess)."""
+
+StateT = TypeVar("StateT")
+"""A kernel state pytree; used where a method returns the same state type it received."""

numpyro/infer/__init__.py

@@ -12,6 +12,7 @@
     TraceMeanField_ELBO,
 )
 from numpyro.infer.ensemble import AIES, ESS
+from numpyro.infer.gibbs import CustomGibbs, DiscreteGibbs, Gibbs
 from numpyro.infer.hmc import HMC, NUTS
 from numpyro.infer.hmc_gibbs import HMCECS, DiscreteHMCGibbs, HMCGibbs
 from numpyro.infer.importance import psis_diagnostic
@@ -50,9 +51,12 @@
     "psis_diagnostic",
     "reparam",
     "BarkerMH",
+    "CustomGibbs",
+    "DiscreteGibbs",
     "DiscreteHMCGibbs",
     "ELBO",
     "ESS",
+    "Gibbs",
     "HMC",
     "HMCECS",
     "HMCGibbs",