diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9ea250d68..5cb3fd961 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/docs/source/mcmc.rst b/docs/source/mcmc.rst index 4c832329f..dc1171511 100644 --- a/docs/source/mcmc.rst +++ b/docs/source/mcmc.rst @@ -9,6 +9,7 @@ We provide a high-level overview of the MCMC algorithms in NumPyro: * `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 `_ combines HMC/NUTS steps with custom Gibbs updates. Gibbs updates must be specified by the user. * `DiscreteHMCGibbs `_ combines HMC/NUTS steps with Gibbs updates for discrete latent variables. The corresponding Gibbs updates are computed automatically. +* `Gibbs `_ composes any number of block kernels (HMC/NUTS, `DiscreteGibbs `_, `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 `_ 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 `_ 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 `_ 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 diff --git a/numpyro/_typing.py b/numpyro/_typing.py index d53d60c17..e3f08ebb8 100644 --- a/numpyro/_typing.py +++ b/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.""" diff --git a/numpyro/infer/__init__.py b/numpyro/infer/__init__.py index a2f4ef264..630482721 100644 --- a/numpyro/infer/__init__.py +++ b/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", diff --git a/numpyro/infer/gibbs.py b/numpyro/infer/gibbs.py new file mode 100644 index 000000000..5c6d33c0a --- /dev/null +++ b/numpyro/infer/gibbs.py @@ -0,0 +1,1224 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +""" +Composable Gibbs kernels: a composite :class:`Gibbs` kernel that updates the latent sites of a +model block by block, the generic block kernels :class:`CustomGibbs` and :class:`DiscreteGibbs`, +and the pure helpers they share with the HMC-within-Gibbs kernels +(:mod:`numpyro.infer.hmc_gibbs`, :mod:`numpyro.infer.mixed_hmc`). +""" + +from collections import OrderedDict +from collections.abc import Callable, Sequence +import copy +from functools import partial, reduce +from typing import Any, NamedTuple, Protocol, TypeAlias + +import numpy as np + +import jax +from jax import random +from jax.flatten_util import ravel_pytree +import jax.numpy as jnp +from jax.scipy.special import expit + +from numpyro._typing import ( + ConstrainFn, + ModelArgs, + ModelKwargs, + ModelT, + PotentialFn, + PyTree, + SiteValues, + TraceT, +) +from numpyro.handlers import condition, seed, substitute, trace +from numpyro.infer.hmc import HMC +from numpyro.infer.initialization import init_to_sample +from numpyro.infer.mcmc import MCMCKernel +from numpyro.infer.util import ( + _prepare_model_for_potential, + _transforms_from_trace, + potential_energy, + transform_fn, +) +from numpyro.util import cond, fori_loop, identity, is_prng_key + +__all__ = [ + "CustomGibbs", + "CustomGibbsState", + "DiscreteGibbs", + "DiscreteGibbsState", + "GIBBS_SITES_KWARG", + "Gibbs", + "GibbsState", + "conditioned", + "discrete_latent_sites", + "with_conditioning", +] + +GIBBS_SITES_KWARG: str = "_gibbs_sites" +"""Reserved model keyword through which conditioning values travel.""" + +ModelWrapper: TypeAlias = Callable[[ModelT], ModelT] +"""Maps a model to a model with the same call signature (conditioning, likelihood estimation).""" + +SiteSelector: TypeAlias = Callable[[TraceT], Sequence[str]] +"""Picks site names from a prototype trace, e.g. :func:`discrete_latent_sites`.""" + +SitesSpec: TypeAlias = Sequence[str] | SiteSelector | None +"""How a block declares its sites: explicit names, a selector, or `None` for the remainder.""" + + +class GibbsUpdateFn(Protocol): + """ + Signature of the user callable of :class:`~numpyro.infer.gibbs.CustomGibbs` / + :class:`~numpyro.infer.hmc_gibbs.HMCGibbs`. Called with keywords only; `hmc_sites` holds + the constrained values of every conditioning site (the name is kept for compatibility). + """ + + def __call__( + self, *, rng_key: jax.Array, gibbs_sites: SiteValues, hmc_sites: SiteValues + ) -> SiteValues: ... + + +def _conditioned_model(model: ModelT, *args: Any, **kwargs: Any) -> Any: + """Module-level target of :func:`conditioned` (kept module-level so kernels pickle).""" + values = kwargs.pop(GIBBS_SITES_KWARG, {}) + with condition(data=values), substitute(data=values): + return model(*args, **kwargs) + + +def conditioned(model: ModelT) -> ModelT: + """ + Return a model that pops :data:`GIBBS_SITES_KWARG` from its keyword arguments and runs + `model` under `condition(data=values)` and `substitute(data=values)`. Idempotent: wrapping + an already conditioned model returns it unchanged. + + :param model: the model. + :return: the conditioned model. + """ + if isinstance(model, partial) and model.func is _conditioned_model: + return model + return partial(_conditioned_model, model) + + +def with_conditioning( + model_kwargs: ModelKwargs | None, + values: SiteValues, + *, + allowed: frozenset[str] | None = None, +) -> ModelKwargs: + """ + Return a copy of `model_kwargs` whose :data:`GIBBS_SITES_KWARG` entry is the existing + entry (if any) extended by `values`. Extending rather than replacing is what makes nested + composites correct: at any depth a block sees exactly the sites it does not own. + + :param model_kwargs: keyword arguments of the model. + :param values: conditioning values to add. + :param allowed: when given, raise if a key of `values` is outside this set. + :return: a new keyword argument dict. + """ + if allowed is not None: + extra = set(values) - allowed + if extra: + raise ValueError(f"Cannot condition on non-latent sites {sorted(extra)}.") + model_kwargs = {} if model_kwargs is None else dict(model_kwargs) + model_kwargs[GIBBS_SITES_KWARG] = { + **model_kwargs.get(GIBBS_SITES_KWARG, {}), + **values, + } + return model_kwargs + + +def prototype_trace( + model: ModelT, + rng_key: jax.Array, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, +) -> TraceT: + """ + Trace `model` once with values drawn by :func:`~numpyro.infer.initialization.init_to_sample` + (which also handles sites without a `sample` method, such as `ImproperUniform`). Callers + must not store the returned trace on a kernel object (it may hold tracers under + `pmap`/`vmap`); store only static metadata derived from it. + + :param model: the model. + :param rng_key: random key used to draw the prototype values. + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model. + :return: the trace. + """ + model_kwargs = {} if model_kwargs is None else model_kwargs + return trace( + substitute(seed(model, rng_key), substitute_fn=init_to_sample) + ).get_trace(*model_args, **model_kwargs) + + +def latent_sample_sites(model_trace: TraceT) -> tuple[str, ...]: + """Names of unobserved `sample` sites in trace order.""" + return tuple( + name + for name, site in model_trace.items() + if site["type"] == "sample" and not site["is_observed"] + ) + + +def discrete_latent_sites(model_trace: TraceT) -> tuple[str, ...]: + """ + Unobserved sample sites whose distribution has enumerate support and which are not + marked `infer={"enumerate": "parallel"}`. Usable as a :data:`SitesSpec` selector. + """ + return tuple( + name + for name, site in model_trace.items() + if site["type"] == "sample" + and not site["is_observed"] + and site["fn"].has_enumerate_support + and site["infer"].get("enumerate", "") != "parallel" + ) + + +def discrete_support_sizes( + model_trace: TraceT, sites: Sequence[str] +) -> dict[str, np.ndarray]: + """Per-site support sizes broadcast to the site's shape, as static `numpy` arrays.""" + return { + name: np.broadcast_to( + model_trace[name]["fn"].enumerate_support(False).shape[0], + jnp.shape(model_trace[name]["value"]), + ) + for name in sites + } + + +def _flat_support_sizes(model_trace: TraceT, sites: Sequence[str]) -> np.ndarray: + """Support sizes flattened in :func:`ravel_pytree` leaf order, as a static numpy array. + + Built with pure numpy so that it stays concrete when `init` runs under a staging + trace such as :func:`jax.pmap` (`jnp` operations would produce tracers there). + """ + sizes = discrete_support_sizes(model_trace, sites) + return np.concatenate([np.ravel(leaf) for leaf in jax.tree.leaves(sizes)]) + + +def subsample_plate_sizes(model_trace: TraceT) -> dict[str, tuple[int, int]]: + """`{plate_name: (size, subsample_size)}` for plates with `size > subsample_size`.""" + return { + name: site["args"] + for name, site in model_trace.items() + if site["type"] == "plate" + and (site["args"][1] is not None) + and site["args"][0] > site["args"][1] + } + + +def any_changed(old: PyTree, new: PyTree) -> jax.Array: + """Scalar boolean: whether any leaf of two pytrees with the same structure differs.""" + flags = [ + jnp.any(a != b) + for a, b in zip(jax.tree.leaves(old), jax.tree.leaves(new), strict=True) + ] + if not flags: + return jnp.array(False) + return reduce(jnp.logical_or, flags) + + +class GibbsState(NamedTuple): + """ + - **z** - dict of the current values of all latent sites, each in the native representation + of the block that owns it (unconstrained for HMC blocks, constrained for + :class:`DiscreteGibbs` and :class:`CustomGibbs` blocks). Written once per step from the + block states; never read back by :meth:`Gibbs.sample`. + - **block_states** - tuple with one state pytree per block, in block order. Source of truth + during a sweep. Addressable from `extra_fields` as `"block_states.."`. + - **rng_key** - random key for the next step. + """ + + z: SiteValues + block_states: tuple[PyTree, ...] + rng_key: jax.Array + + +def _as_arrays(values: SiteValues) -> SiteValues: + """Canonicalize values to arrays so that `cond` branches carrying them agree on dtypes.""" + return {k: jnp.asarray(v) for k, v in values.items()} + + +def _or(flags: Sequence[Any]) -> Any: + """ + Logical or of Python bools and traced booleans. Returns a Python bool when the result is + statically known, so that callers can skip emitting a `cond`. + """ + traced = [f for f in flags if not isinstance(f, bool)] + if any(f is True for f in flags): + return True + if not traced: + return False + return reduce(jnp.logical_or, traced) + + +def _maybe_refresh( + kernel: MCMCKernel, + state: PyTree, + changed: Any, + model_args: ModelArgs, + model_kwargs: ModelKwargs, +) -> PyTree: + if changed is False: + return state + if changed is True: + return kernel.refresh(state, model_args, model_kwargs) + return cond( + changed, + state, + lambda s: kernel.refresh(s, model_args, model_kwargs), + state, + identity, + ) + + +def _has_model(kernel: MCMCKernel) -> bool: + return getattr(kernel, "model", None) is not None + + +class Gibbs(MCMCKernel): + """ + Composite kernel that updates the latent sites of a model block by block. Each block is a + `(kernel, sites)` pair; the block kernel is run on the model conditioned on the current + values of every site it does not own. Blocks are visited in order once per MCMC step. + + :param blocks: sequence of `(kernel, sites)`. `kernel` is any + :class:`~numpyro.infer.mcmc.MCMCKernel` that implements + :meth:`~numpyro.infer.mcmc.MCMCKernel.refresh` (and + :meth:`~numpyro.infer.mcmc.MCMCKernel.wrap_model` if it holds a model). `sites` is a + sequence of site names, a callable mapping a prototype trace to site names (for example + :func:`~numpyro.infer.gibbs.discrete_latent_sites`), or `None` for "all latent + sample sites not owned by another block" (allowed for at most one block). All + model-based blocks must be built on the same model callable. + + Validation at construction: at least one block; kernels override `refresh`; model-based + kernels share one model; each model-based kernel is rebound once with + ``kernel.wrap_model(conditioned)``. + + Validation at `init` (errors, never warnings, because the fallbacks inside + :func:`~numpyro.infer.util.initialize_model` are silent): the union of block sites equals + the set of unobserved sample sites minus those marked `enumerate="parallel"`; no discrete + latent is left to an HMC block; blocks are disjoint and non-empty; no site is already + conditioned by an enclosing composite; `init_params` keys lie in the union (unconstrained + values for HMC blocks, constrained values otherwise); the model has no subsample plates + (use :class:`~numpyro.infer.hmc_gibbs.HMCECS`); HMC blocks have no value-dependent + supports. + + .. note:: Each HMC block adapts its step size and mass matrix on its own conditional, so + blocks should leave `find_heuristic_step_size=False` (the heuristic binds the potential + to the initial conditioning). Strongly correlated continuous sites belong in one HMC + block. + + **Example** + + .. doctest:: + + >>> from jax import random + >>> import jax.numpy as jnp + >>> import numpyro + >>> import numpyro.distributions as dist + >>> from numpyro.infer import MCMC, NUTS, Gibbs, DiscreteGibbs, CustomGibbs + >>> from numpyro.infer.gibbs import discrete_latent_sites + ... + >>> def model(probs, locs): + ... c = numpyro.sample("c", dist.Categorical(probs)) + ... x = numpyro.sample("x", dist.Normal(locs[c], 0.5)) + ... y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + ... numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + ... + >>> def gibbs_fn(rng_key, gibbs_sites, hmc_sites): + ... x = hmc_sites["x"] + ... return {"y": dist.Normal(0.8 * (1 - x), jnp.sqrt(0.8)).sample(rng_key)} + ... + >>> kernel = Gibbs([ + ... (DiscreteGibbs(model), discrete_latent_sites), + ... (CustomGibbs(gibbs_fn), ["y"]), + ... (NUTS(model), None), + ... ]) + >>> mcmc = MCMC(kernel, num_warmup=100, num_samples=100, progress_bar=False) + >>> mcmc.run(random.key(0), jnp.array([0.15, 0.3, 0.3, 0.25]), jnp.array([-2.0, 0.0, 2.0, 4.0])) + >>> mcmc.print_summary() # doctest: +SKIP + """ + + _state_cls: type[GibbsState] = GibbsState + sample_field: str = "z" + + def __init__(self, blocks: Sequence[tuple[MCMCKernel, SitesSpec]]) -> None: + blocks = list(blocks) + if not blocks: + raise ValueError("Gibbs requires at least one block.") + kernels, specs = [], [] + for block in blocks: + if not (isinstance(block, tuple) and len(block) == 2): + raise ValueError("Each block must be a `(kernel, sites)` pair.") + kernel, spec = block + if not isinstance(kernel, MCMCKernel): + raise ValueError(f"{kernel!r} is not an MCMCKernel.") + if type(kernel).refresh is MCMCKernel.refresh: + raise ValueError( + f"{type(kernel).__name__} does not implement `refresh` and cannot be " + "used as a block of Gibbs." + ) + if hasattr(kernel, "model") and kernel.model is None: + raise ValueError( + "Kernels built from a potential function cannot be blocks of Gibbs." + ) + if spec is not None and not callable(spec): + spec = tuple(spec) + if not spec or not all(isinstance(s, str) for s in spec): + raise ValueError( + "`sites` must be a non-empty sequence of site names, a callable " + "or None." + ) + kernels.append(kernel) + specs.append(spec) + if sum(spec is None for spec in specs) > 1: + raise ValueError("At most one block can use `None` for its sites.") + models = [kernel.model for kernel in kernels if _has_model(kernel)] + if not models: + raise ValueError("Gibbs requires at least one block built on a model.") + if any(model != models[0] for model in models[1:]): + raise ValueError("All model-based blocks must share the same model.") + self._model = models[0] + self._kernels = tuple( + kernel.wrap_model(conditioned) if _has_model(kernel) else kernel + for kernel in kernels + ) + self._specs = tuple(specs) + # static metadata resolved at `init` + self._sites: tuple[tuple[str, ...], ...] = () + self._has_deterministic = False + self._sample_fn = None + + @property + def model(self) -> ModelT | None: + """The shared, unwrapped model of the model-based blocks.""" + return self._model + + @property + def blocks(self) -> tuple[tuple[MCMCKernel, SitesSpec], ...]: + """The `(kernel, sites)` pairs; kernels are the conditioned copies actually run.""" + return tuple(zip(self._kernels, self._specs)) + + @property + def default_fields(self) -> tuple[str, ...]: + return ("z",) + + def get_diagnostics_str(self, state: GibbsState) -> str: + parts = [ + kernel.get_diagnostics_str(block_state) + for kernel, block_state in zip(self._kernels, state.block_states) + ] + return " | ".join(part for part in parts if part) + + def _resolve_partition( + self, model_trace: Any, enclosing: frozenset[str] + ) -> tuple[tuple[str, ...], ...]: + latent = tuple( + name + for name in latent_sample_sites(model_trace) + if model_trace[name]["infer"].get("enumerate", "") != "parallel" + ) + latent_set = frozenset(latent) + sites: list[tuple[str, ...] | None] = [] + for spec in self._specs: + if spec is None: + sites.append(None) + continue + names = tuple(spec(model_trace)) if callable(spec) else spec + unknown = [name for name in names if name not in latent_set] + if unknown: + hint = ( + " (already conditioned by an enclosing kernel)" + if any(name in enclosing for name in unknown) + else "" + ) + raise ValueError( + f"Sites {unknown} are not latent sample sites of the model{hint}." + ) + sites.append(names) + owned = [name for names in sites if names is not None for name in names] + duplicates = sorted({name for name in owned if owned.count(name) > 1}) + if duplicates: + raise ValueError(f"Sites {duplicates} are owned by more than one block.") + remainder = tuple(name for name in latent if name not in owned) + # the remainder block may own no sites (for example an HMC block on a model whose + # latent sites are all discrete); explicit blocks may not + if None in sites: + sites[sites.index(None)] = remainder + elif remainder: + raise ValueError( + f"Latent sites {list(remainder)} are not owned by any block; add a block " + "or use `None` as the sites of one block." + ) + resolved = tuple(names if names is not None else () for names in sites) + for kernel, spec, names in zip(self._kernels, self._specs, resolved): + if not names and spec is not None: + raise ValueError(f"Block {type(kernel).__name__} owns no sites.") + if isinstance(kernel, HMC): + discrete = [ + name + for name in names + if model_trace[name]["fn"].support.is_discrete + ] + if discrete: + raise ValueError( + f"Discrete latent sites {discrete} cannot be sampled by an HMC " + "block; use DiscreteGibbs or mark them with " + "`infer={'enumerate': 'parallel'}`." + ) + return resolved + + @staticmethod + def _block_trace(model_trace: Any, names: Sequence[str]) -> Any: + """The trace restricted to the block's sample sites (other site types are kept).""" + return OrderedDict( + (name, site) + for name, site in model_trace.items() + if site["type"] != "sample" or name in names + ) + + def _split_init_params( + self, init_params: SiteValues | None + ) -> tuple[SiteValues | None, ...]: + if not init_params: + return tuple(None for _ in self._kernels) + owned = {name for names in self._sites for name in names} + unknown = sorted(set(init_params) - owned) + if unknown: + raise ValueError(f"`init_params` has unknown sites {unknown}.") + return tuple( + {name: init_params[name] for name in names if name in init_params} or None + for names in self._sites + ) + + def init( + self, + rng_key: jax.Array, + num_warmup: int, + init_params: SiteValues | None, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> GibbsState: + if not is_prng_key(rng_key): + raise ValueError( + "Gibbs only supports a single random key; for multiple chains use " + '`chain_method="parallel"`, `chain_method="sequential"` or a callable ' + "chain method such as `jax.vmap`." + ) + model_kwargs = {} if model_kwargs is None else dict(model_kwargs) + enclosing = frozenset(model_kwargs.get(GIBBS_SITES_KWARG, {})) + rng_key, key_trace = random.split(rng_key) + model_trace = prototype_trace( + conditioned(self._model), key_trace, model_args, model_kwargs + ) + if subsample_plate_sizes(model_trace): + raise ValueError( + "Gibbs does not support models with subsample plates; use HMCECS." + ) + self._sites = self._resolve_partition(model_trace, enclosing) + self._has_deterministic = any( + site["type"] == "deterministic" for site in model_trace.values() + ) + block_init_params = self._split_init_params(init_params) + + # constrained initial values of every latent site, used to condition the blocks + constrained = { + name: model_trace[name]["value"] for names in self._sites for name in names + } + for kernel, names, params in zip(self._kernels, self._sites, block_init_params): + if isinstance(kernel, HMC): + transforms = _transforms_from_trace( + self._block_trace(model_trace, names), raise_warnings=False + ) + if transforms.dynamic_support: + raise ValueError( + f"The supports of sites {list(names)} depend on other sites; " + "Gibbs does not support value-dependent supports across blocks." + ) + if params: + constrained.update(transform_fn(transforms.inv_transforms, params)) + elif params: + constrained.update(params) + + block_states = [] + for i, (kernel, names, params) in enumerate( + zip(self._kernels, self._sites, block_init_params) + ): + siblings = {k: v for k, v in constrained.items() if k not in names} + kwargs_i = with_conditioning(model_kwargs, siblings) + if not isinstance(kernel, HMC): + params = {**{k: constrained[k] for k in names}, **(params or {})} + rng_key, key_i = random.split(rng_key) + state_i = kernel.init(key_i, num_warmup, params, model_args, kwargs_i) + block_states.append(state_i) + z_i = getattr(state_i, kernel.sample_field) + constrained_i = kernel.get_constrain_fn(model_args, kwargs_i)(z_i) + constrained.update({k: constrained_i[k] for k in names}) + + # blocks were initialized against prototype values of later blocks; refresh them + for i in range(len(self._kernels) - 1): + siblings = {k: v for k, v in constrained.items() if k not in self._sites[i]} + kwargs_i = with_conditioning(model_kwargs, siblings) + block_states[i] = self._kernels[i].refresh( + block_states[i], model_args, kwargs_i + ) + + z = self._merge_z(block_states) + self._sample_fn = self._sample_one + return self._state_cls(z, tuple(block_states), rng_key) + + def _merge_z(self, block_states: Sequence[PyTree]) -> SiteValues: + z = {} + for kernel, names, block_state in zip(self._kernels, self._sites, block_states): + z_i = getattr(block_state, kernel.sample_field) + z.update({k: z_i[k] for k in names}) + return z + + def _sample_one( + self, + state: GibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> GibbsState: + model_kwargs = {} if model_kwargs is None else model_kwargs + num_blocks = len(self._kernels) + block_states = list(state.block_states) + z_blocks = [ + getattr(block_state, kernel.sample_field) + for kernel, block_state in zip(self._kernels, block_states) + ] + # constrained values of each block's own sites, for conditioning its siblings + constrained = [ + { + k: v + for k, v in kernel.get_constrain_fn(model_args, model_kwargs)( + z_i + ).items() + if k in names + } + for kernel, names, z_i in zip(self._kernels, self._sites, z_blocks) + ] + + def kwargs_for(i: int) -> ModelKwargs: + siblings = {} + for j in range(num_blocks): + if j != i: + siblings.update(constrained[j]) + return with_conditioning(model_kwargs, siblings) + + moved: list[Any] = [False] * num_blocks + for i, (kernel, names) in enumerate(zip(self._kernels, self._sites)): + kwargs_i = kwargs_for(i) + # siblings visited earlier in this sweep may have moved since the last refresh + block_states[i] = _maybe_refresh( + kernel, block_states[i], _or(moved[:i]), model_args, kwargs_i + ) + block_states[i] = kernel.sample(block_states[i], model_args, kwargs_i) + z_new = getattr(block_states[i], kernel.sample_field) + moved[i] = any_changed(z_blocks[i], z_new) + constrained[i] = { + k: v + for k, v in kernel.get_constrain_fn(model_args, kwargs_i)(z_new).items() + if k in names + } + # siblings visited later in this sweep may have moved; refresh against the final values + for i in range(num_blocks - 1): + block_states[i] = _maybe_refresh( + self._kernels[i], + block_states[i], + _or(moved[i + 1 :]), + model_args, + kwargs_for(i), + ) + + rng_key, _ = random.split(state.rng_key) + return state._replace( + z=self._merge_z(block_states), + block_states=tuple(block_states), + rng_key=rng_key, + ) + + def sample( + self, + state: GibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> GibbsState: + """ + Run one sweep over the blocks from the given :class:`GibbsState` and return the + resulting :class:`GibbsState`. + + :param GibbsState state: the current state. + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model. + :return: the next state. + """ + assert self._sample_fn is not None, "`init` must be called before `sample`." + return self._sample_fn(state, model_args, model_kwargs) + + def refresh( + self, + state: GibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> GibbsState: + """Identity: :meth:`sample` refreshes each block against the current conditioning.""" + return state + + def wrap_model(self, wrapper: ModelWrapper) -> "Gibbs": + """New composite with `wrapper` applied to every model-based block (nesting).""" + kernel = copy.copy(self) + kernel._model = wrapper(self._model) + kernel._kernels = tuple( + block.wrap_model(wrapper) if _has_model(block) else block + for block in self._kernels + ) + kernel._sites = () + kernel._sample_fn = None + return kernel + + def get_constrain_fn( + self, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> ConstrainFn: + """Constrain each block's sites with the block's own constrain function (no replay).""" + if not self._sites: + return identity + + def fn(z: SiteValues) -> SiteValues: + out = {} + for kernel, names in zip(self._kernels, self._sites): + z_i = {k: z[k] for k in names} + constrained_i = kernel.get_constrain_fn(model_args, model_kwargs)(z_i) + out.update({k: constrained_i[k] for k in names}) + return out + + return fn + + def postprocess_fn( + self, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> ConstrainFn: + """ + Constrain each block's sites with the block's own constrain function, then, only if + the model has `deterministic` sites, replay the model once with all constrained values + substituted to collect them. + """ + if not self._sites: + return identity + constrain = self.get_constrain_fn(model_args, model_kwargs) + if not self._has_deterministic: + return constrain + model_kwargs = {} if model_kwargs is None else model_kwargs + + def fn(z: SiteValues) -> SiteValues: + constrained = constrain(z) + model = substitute( + seed(conditioned(self._model), random.key(0)), data=constrained + ) + model_trace = trace(model).get_trace(*model_args, **model_kwargs) + deterministic = { + name: site["value"] + for name, site in model_trace.items() + if site["type"] == "deterministic" + } + return {**constrained, **deterministic} + + return fn + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state["_sample_fn"] = None + return state + + +class CustomGibbsState(NamedTuple): + """ + - **z** - dict of the block's current (constrained) values. + - **rng_key** - random key for the next step. + """ + + z: SiteValues + rng_key: jax.Array + + +class CustomGibbs(MCMCKernel): + """ + Block kernel that delegates the update to a user callable. The callable receives the + current values of the block's sites and the constrained values of all conditioning sites + and returns new values for the block's sites (it must sample from the conditional; + correctness is the user's responsibility, as with + :class:`~numpyro.infer.hmc_gibbs.HMCGibbs`). Only usable as a block of :class:`Gibbs`. + + :param gibbs_fn: called as `gibbs_fn(rng_key=..., gibbs_sites=..., hmc_sites=...)`; the + keyword names are kept for compatibility with existing `HMCGibbs` users, and + `hmc_sites` holds all conditioning sites (not only HMC ones). The returned dict must + have exactly the block's site names as keys. + """ + + sample_field: str = "z" + + def __init__(self, gibbs_fn: GibbsUpdateFn) -> None: + if not callable(gibbs_fn): + raise ValueError("gibbs_fn must be a callable") + self._gibbs_fn = gibbs_fn + + def init( + self, + rng_key: jax.Array, + num_warmup: int, + init_params: SiteValues | None, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> CustomGibbsState: + """`init_params` are the block's initial (constrained) values, supplied by the composite.""" + if not init_params: + raise ValueError( + "CustomGibbs requires initial values; use it as a block of Gibbs." + ) + return CustomGibbsState(_as_arrays(init_params), rng_key) + + def sample( + self, + state: CustomGibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> CustomGibbsState: + """Reads the conditioning values from `model_kwargs[GIBBS_SITES_KWARG]`.""" + model_kwargs = {} if model_kwargs is None else model_kwargs + rng_key, key_update = random.split(state.rng_key) + z_new = self._gibbs_fn( + rng_key=key_update, + gibbs_sites=state.z, + hmc_sites=model_kwargs.get(GIBBS_SITES_KWARG, {}), + ) + if set(z_new) != set(state.z): + raise ValueError( + f"gibbs_fn must return values for exactly the sites {sorted(state.z)}, " + f"got {sorted(z_new)}." + ) + return state._replace(z=_as_arrays(z_new), rng_key=rng_key) + + def refresh( + self, + state: CustomGibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> CustomGibbsState: + """Identity: nothing is cached.""" + return state + + def wrap_model(self, wrapper: ModelWrapper) -> "CustomGibbs": + """Returns `self`: there is no model.""" + return self + + def get_constrain_fn( + self, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> Callable: + """Identity: the block's values are already constrained.""" + return identity + + +# Discrete proposals. Each is called as +# `(rng_key, z, pe, potential_fn, idx, support_size) -> (rng_key, z_new, pe_new, log_accept_ratio)` +# where `idx` is the flat coordinate of `z` to update and `support_size` its support size. + +ProposalFn: TypeAlias = Callable[ + [jax.Array, SiteValues, jax.Array, PotentialFn, jax.Array, jax.Array], + tuple[jax.Array, SiteValues, jax.Array, jax.Array], +] +"""`(rng_key, z, pe, potential_fn, idx, support_size) -> (rng_key, z_new, pe_new, log_accept_ratio)`.""" + + +def _discrete_gibbs_proposal_body_fn( + z_init_flat, unravel_fn, pe_init, potential_fn, idx, i, val +): + rng_key, z, pe, log_weight_sum = val + rng_key, rng_transition = random.split(rng_key) + proposal = jnp.where(i >= z_init_flat[idx], i + 1, i) + z_new_flat = z_init_flat.at[idx].set(proposal) + z_new = unravel_fn(z_new_flat) + pe_new = potential_fn(z_new) + log_weight_new = pe_init - pe_new + # Handles the NaN case... + log_weight_new = jnp.where(jnp.isfinite(log_weight_new), log_weight_new, -jnp.inf) + # transition_prob = e^weight_new / (e^weight_logsumexp + e^weight_new) + transition_prob = expit(log_weight_new - log_weight_sum) + z, pe = cond( + random.bernoulli(rng_transition, transition_prob), + (z_new, pe_new), + identity, + (z, pe), + identity, + ) + log_weight_sum = jnp.logaddexp(log_weight_new, log_weight_sum) + return rng_key, z, pe, log_weight_sum + + +def _discrete_gibbs_proposal( + rng_key: jax.Array, + z_discrete: SiteValues, + pe: jax.Array, + potential_fn: PotentialFn, + idx: jax.Array, + support_size: jax.Array, +) -> tuple[jax.Array, SiteValues, jax.Array, jax.Array]: + # idx: current index of `z_discrete_flat` to update + # support_size: support size of z_discrete at the index idx + + z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) + # Here we loop over the support of z_flat[idx] to get z_new + # Note: we can't vmap potential_fn over all proposals and sample from the conditional + # categorical distribution because support_size is a traced value, i.e. its value + # might change across different discrete variables; + # so here we will loop over all proposals and use an online scheme to sample from + # the conditional categorical distribution + body_fn = partial( + _discrete_gibbs_proposal_body_fn, + z_discrete_flat, + unravel_fn, + pe, + potential_fn, + idx, + ) + init_val = (rng_key, z_discrete, pe, jnp.array(0.0)) + rng_key, z_new, pe_new, _ = fori_loop(0, support_size - 1, body_fn, init_val) + log_accept_ratio = jnp.array(0.0) + return rng_key, z_new, pe_new, log_accept_ratio + + +def _discrete_modified_gibbs_proposal( + rng_key: jax.Array, + z_discrete: SiteValues, + pe: jax.Array, + potential_fn: PotentialFn, + idx: jax.Array, + support_size: jax.Array, + stay_prob: float = 0.0, +) -> tuple[jax.Array, SiteValues, jax.Array, jax.Array]: + assert isinstance(stay_prob, float) and stay_prob >= 0.0 and stay_prob < 1 + z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) + body_fn = partial( + _discrete_gibbs_proposal_body_fn, + z_discrete_flat, + unravel_fn, + pe, + potential_fn, + idx, + ) + # like gibbs_step but here, weight of the current value is 0 + init_val = (rng_key, z_discrete, pe, jnp.array(-jnp.inf)) + rng_key, z_new, pe_new, log_weight_sum = fori_loop( + 0, support_size - 1, body_fn, init_val + ) + rng_key, rng_stay = random.split(rng_key) + z_new, pe_new = cond( + random.bernoulli(rng_stay, stay_prob), + (z_discrete, pe), + identity, + (z_new, pe_new), + identity, + ) + # here we calculate the MH correction: (1 - P(z)) / (1 - P(z_new)) + # where 1 - P(z) ~ weight_sum + # and 1 - P(z_new) ~ 1 + weight_sum - z_new_weight + log_accept_ratio = log_weight_sum - jnp.log( + jnp.exp(log_weight_sum) - jnp.expm1(pe - pe_new) + ) + return rng_key, z_new, pe_new, log_accept_ratio + + +def _discrete_rw_proposal( + rng_key: jax.Array, + z_discrete: SiteValues, + pe: jax.Array, + potential_fn: PotentialFn, + idx: jax.Array, + support_size: jax.Array, +) -> tuple[jax.Array, SiteValues, jax.Array, jax.Array]: + rng_key, rng_proposal = random.split(rng_key, 2) + z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) + + proposal = random.randint(rng_proposal, (), minval=0, maxval=support_size) + z_new_flat = z_discrete_flat.at[idx].set(proposal) + z_new = unravel_fn(z_new_flat) + pe_new = potential_fn(z_new) + log_accept_ratio = pe - pe_new + return rng_key, z_new, pe_new, log_accept_ratio + + +def _discrete_modified_rw_proposal( + rng_key: jax.Array, + z_discrete: SiteValues, + pe: jax.Array, + potential_fn: PotentialFn, + idx: jax.Array, + support_size: jax.Array, + stay_prob: float = 0.0, +) -> tuple[jax.Array, SiteValues, jax.Array, jax.Array]: + assert isinstance(stay_prob, float) and stay_prob >= 0.0 and stay_prob < 1 + rng_key, rng_proposal, rng_stay = random.split(rng_key, 3) + z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) + + i = random.randint(rng_proposal, (), minval=0, maxval=support_size - 1) + proposal = jnp.where(i >= z_discrete_flat[idx], i + 1, i) + proposal = jnp.where( + random.bernoulli(rng_stay, stay_prob), z_discrete_flat[idx], proposal + ) + z_new_flat = z_discrete_flat.at[idx].set(proposal) + z_new = unravel_fn(z_new_flat) + pe_new = potential_fn(z_new) + log_accept_ratio = pe - pe_new + return rng_key, z_new, pe_new, log_accept_ratio + + +def select_discrete_proposal(random_walk: bool, modified: bool) -> ProposalFn: + """ + Pick the discrete proposal: the exact conditional (Gibbs) or a uniform random walk, each + optionally in Liu's modified form that never proposes the current value. + + :param bool random_walk: uniform proposals over the support instead of the conditional. + :param bool modified: use the modified (Metropolised) proposal. + """ + if random_walk: + if modified: + return partial(_discrete_modified_rw_proposal, stay_prob=0.0) + return _discrete_rw_proposal + if modified: + return partial(_discrete_modified_gibbs_proposal, stay_prob=0.0) + return _discrete_gibbs_proposal + + +def discrete_gibbs_sweep( + rng_key: jax.Array, + z: SiteValues, + potential_energy: jax.Array, + potential_fn: PotentialFn, + support_sizes_flat: jax.Array, + proposal_fn: ProposalFn, +) -> tuple[SiteValues, jax.Array]: + """ + One sweep over the flat discrete coordinates of `z` in a random order, each coordinate + updated with `proposal_fn` and Metropolis corrected. + + :param rng_key: random key. + :param z: current discrete values. + :param potential_energy: potential energy at `z`. + :param potential_fn: potential energy as a function of the discrete values. + :param support_sizes_flat: support size of each flat coordinate, in `ravel_pytree` order. + :param proposal_fn: a discrete proposal, see :func:`select_discrete_proposal`. + :return: the new values and their potential energy. + """ + num_discretes = support_sizes_flat.shape[0] + rng_key, rng_permute = random.split(rng_key) + idxs = random.permutation(rng_permute, jnp.arange(num_discretes)) + + def body_fn(i, val): + idx = idxs[i] + support_size = support_sizes_flat[idx] + rng_key, z, pe = val + rng_key, z_new, pe_new, log_accept_ratio = proposal_fn( + rng_key, z, pe, potential_fn, idx, support_size + ) + rng_key, rng_accept = random.split(rng_key) + # u ~ Uniform(0, 1), u < accept_ratio => -log(u) > -log_accept_ratio + # and -log(u) ~ exponential(1) + z, pe = cond( + random.exponential(rng_accept) > -log_accept_ratio, + (z_new, pe_new), + identity, + (z, pe), + identity, + ) + return rng_key, z, pe + + init_val = (rng_key, z, potential_energy) + _, z, pe = fori_loop(0, num_discretes, body_fn, init_val) + return z, pe + + +class DiscreteGibbsState(NamedTuple): + """ + - **z** - dict of the current discrete values. + - **potential_energy** - potential energy at `z` under the current conditioning (the block's + own potential, which differs from an HMC block's by a constant and cannot be handed + across blocks). + - **rng_key** - random key for the next step. + """ + + z: SiteValues + potential_energy: jax.Array + rng_key: jax.Array + + +class DiscreteGibbs(MCMCKernel): + """ + Metropolis / Gibbs updates of the discrete latent sites of a model, one flat coordinate at + a time in a random order. Usable standalone on a purely discrete model, or as a block of + :class:`Gibbs`. + + The potential is the negative log joint of the model at the block's (constrained) values, + with every other latent site fixed to the conditioning values passed through the model + keyword arguments; sites marked `infer={"enumerate": "parallel"}` are marginalized exactly + as in HMC. + + :param model: the model. + :param bool random_walk: uniform proposals over the support instead of the exact + conditional. + :param bool modified: Liu's modified (Metropolised) proposal that never proposes the + current value. + + **References:** + + 1. *Peskun's theorem and a modified discrete-state Gibbs sampler*, Liu, J. S. (1996) + """ + + sample_field: str = "z" + + def __init__( + self, model: ModelT, *, random_walk: bool = False, modified: bool = False + ) -> None: + self._model = model + self._random_walk = random_walk + self._modified = modified + self._proposal_fn = select_discrete_proposal(random_walk, modified) + # static metadata resolved at `init` + self._sites: tuple[str, ...] | None = None + self._support_sizes_flat: np.ndarray | None = None + self._enum = False + # closes over trace values, rebuilt at every `init` + self._prepared_model: ModelT | None = None + + @property + def model(self) -> ModelT: + return self._model + + def get_potential_fn( + self, + model_args: ModelArgs = (), + model_kwargs: ModelKwargs | None = None, + ) -> PotentialFn: + """Potential over the block's discrete values for the given arguments and conditioning.""" + if self._prepared_model is None: + raise RuntimeError( + "`get_potential_fn` requires the kernel to be initialized; run `init` first." + ) + prepared_model, enum = self._prepared_model, self._enum + + def potential_fn(z: SiteValues) -> jax.Array: + return potential_energy( + prepared_model, + model_args, + with_conditioning(model_kwargs, z), + {}, + enum=enum, + ) + + return potential_fn + + def init( + self, + rng_key: jax.Array, + num_warmup: int, + init_params: SiteValues | None, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> DiscreteGibbsState: + model_kwargs = {} if model_kwargs is None else dict(model_kwargs) + rng_key, key_trace = random.split(rng_key) + model_trace = prototype_trace( + conditioned(self._model), key_trace, model_args, model_kwargs + ) + sites = discrete_latent_sites(model_trace) + if not sites: + raise ValueError( + "Cannot detect any discrete latent variables in the model." + ) + others = [ + name + for name in latent_sample_sites(model_trace) + if name not in sites + and model_trace[name]["infer"].get("enumerate", "") != "parallel" + ] + if others: + raise ValueError( + f"DiscreteGibbs cannot sample the latent sites {others}; condition on them " + "or use DiscreteGibbs as a block of Gibbs." + ) + self._sites = sites + self._support_sizes_flat = _flat_support_sizes(model_trace, sites) + self._enum = any( + site["type"] == "sample" + and not site["is_observed"] + and site["infer"].get("enumerate", "") == "parallel" + for site in model_trace.values() + ) + self._prepared_model = _prepare_model_for_potential( + conditioned(self._model), model_trace, enum=self._enum + ) + init_params = {} if init_params is None else init_params + z = _as_arrays( + {name: init_params.get(name, model_trace[name]["value"]) for name in sites} + ) + pe = self.get_potential_fn(model_args, model_kwargs)(z) + return DiscreteGibbsState(z, jnp.asarray(pe), rng_key) + + def sample( + self, + state: DiscreteGibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> DiscreteGibbsState: + """One :func:`~numpyro.infer.gibbs.discrete_gibbs_sweep` with the selected proposal.""" + rng_key, key_sweep = random.split(state.rng_key) + z, pe = discrete_gibbs_sweep( + key_sweep, + state.z, + state.potential_energy, + self.get_potential_fn(model_args, model_kwargs), + jnp.asarray(self._support_sizes_flat), + self._proposal_fn, + ) + return state._replace(z=z, potential_energy=pe, rng_key=rng_key) + + def refresh( + self, + state: DiscreteGibbsState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> DiscreteGibbsState: + """Recompute `potential_energy` at `state.z` (one model evaluation).""" + pe = self.get_potential_fn(model_args, model_kwargs)(state.z) + return state._replace(potential_energy=jnp.asarray(pe)) + + def wrap_model(self, wrapper: ModelWrapper) -> "DiscreteGibbs": + kernel = copy.copy(self) + kernel._model = wrapper(self._model) + kernel._sites = None + kernel._support_sizes_flat = None + kernel._prepared_model = None + return kernel + + def get_constrain_fn( + self, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> Callable: + """Identity: discrete values are already constrained.""" + return identity + + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state["_prepared_model"] = None + return state diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py index d15b290e8..8dbb60cf1 100644 --- a/numpyro/infer/hmc.py +++ b/numpyro/infer/hmc.py @@ -2,17 +2,23 @@ # SPDX-License-Identifier: Apache-2.0 from collections import OrderedDict, namedtuple +from collections.abc import Callable +import copy from functools import partial import math import os +from typing import cast import warnings +import jax from jax import lax, random, vmap from jax.flatten_util import ravel_pytree import jax.numpy as jnp +from numpyro._typing import ConstrainFn, ModelArgs, ModelKwargs, PotentialFn from numpyro.infer.hmc_util import ( IntegratorState, + _value_and_grad, build_tree, euclidean_kinetic_energy, find_reasonable_step_size, @@ -22,9 +28,11 @@ from numpyro.infer.mcmc import MCMCKernel from numpyro.infer.util import ( ParamInfo, + _transforms_from_trace, find_stack_level, init_to_uniform, initialize_model, + transform_fn, ) from numpyro.util import cond, fori_loop, identity, is_prng_key @@ -89,11 +97,14 @@ def _get_num_steps(step_size, trajectory_length): return num_steps.astype(jnp.result_type(int)) -def momentum_generator(prototype_r, mass_matrix_sqrt, rng_key): +def momentum_generator( + prototype_r, mass_matrix_sqrt: dict[tuple[str, ...], jax.Array] | jax.Array, rng_key +): if isinstance(mass_matrix_sqrt, dict): - rng_keys = random.split(rng_key, len(mass_matrix_sqrt)) + blocks = cast(dict[tuple[str, ...], jax.Array], mass_matrix_sqrt) + rng_keys = random.split(rng_key, len(blocks)) r = {} - for (site_names, mm_sqrt), rng_key in zip(mass_matrix_sqrt.items(), rng_keys): + for (site_names, mm_sqrt), rng_key in zip(blocks.items(), rng_keys): r_block = OrderedDict([(k, prototype_r[k]) for k in site_names]) r.update(momentum_generator(r_block, mm_sqrt, rng_key)) return r @@ -333,7 +344,7 @@ def init_kernel( ) rng_key_hmc, rng_key_wa, rng_key_momentum = random.split(rng_key, 3) - z_info = IntegratorState(z=z, potential_energy=pe, z_grad=z_grad) + z_info = IntegratorState(z=z, r=None, potential_energy=pe, z_grad=z_grad) wa_state = wa_init( z_info, rng_key_wa, step_size, inverse_mass_matrix=inverse_mass_matrix ) @@ -374,6 +385,7 @@ def _hmc_next( pe_fn = potential_fn_gen(*model_args, **model_kwargs) _, vv_update_fn = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad) else: + assert vv_update is not None vv_update_fn = vv_update if fixed_num_steps is not None: @@ -426,8 +438,10 @@ def _nuts_next( pe_fn = potential_fn_gen(*model_args, **model_kwargs) _, vv_update_fn = velocity_verlet(pe_fn, kinetic_fn, forward_mode_ad) else: + assert vv_update is not None vv_update_fn = vv_update + assert max_treedepth is not None binary_tree = build_tree( vv_update_fn, kinetic_fn, @@ -485,6 +499,7 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None): if algo == "HMC": hmc_length_args = (hmc_state.trajectory_length,) else: + assert max_treedepth is not None hmc_length_args = ( jnp.where(hmc_state.i < wa_steps, max_treedepth[0], max_treedepth[1]), ) @@ -498,10 +513,12 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None): *hmc_length_args, ) # not update adapt_state after warmup phase + wa_update_fn = wa_update + assert wa_update_fn is not None adapt_state = cond( hmc_state.i < wa_steps, (hmc_state.i, accept_prob, vv_state, hmc_state.adapt_state), - lambda args: wa_update(*args), + lambda args: wa_update_fn(*args), hmc_state.adapt_state, identity, ) @@ -532,8 +549,8 @@ def sample_kernel(hmc_state, model_args=(), model_kwargs=None): # Make `init_kernel` and `sample_kernel` visible from the global scope once # `hmc` is called for sphinx doc generation. if "SPHINX_BUILD" in os.environ: - hmc.init_kernel = init_kernel - hmc.sample_kernel = sample_kernel + hmc.init_kernel = init_kernel # ty: ignore[unresolved-attribute] + hmc.sample_kernel = sample_kernel # ty: ignore[unresolved-attribute] return init_kernel, sample_kernel @@ -684,6 +701,8 @@ def __init__( self._potential_fn_gen = None self._postprocess_fn = None self._sample_fn = None + self._inv_transforms = None + self._dynamic_support = None def _init_state(self, rng_key, model_args, model_kwargs, init_params): if self._model is not None: @@ -711,6 +730,9 @@ def _init_state(self, rng_key, model_args, model_kwargs, init_params): ) self._potential_fn_gen = potential_fn self._postprocess_fn = postprocess_fn + transforms = _transforms_from_trace(model_trace, raise_warnings=False) + self._inv_transforms = transforms.inv_transforms + self._dynamic_support = transforms.dynamic_support elif self._init_fn is None: self._init_fn, self._sample_fn = hmc( potential_fn=self._potential_fn, @@ -737,6 +759,102 @@ def get_diagnostics_str(self, state): state.num_steps, state.adapt_state.step_size, state.mean_accept_prob ) + def get_potential_fn( + self, + model_args: ModelArgs = (), + model_kwargs: ModelKwargs | None = None, + ) -> PotentialFn: + """ + Return the potential energy function (negative log joint in unconstrained space) for + the given model arguments. Requires :meth:`init` to have run. + + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model. + :return: a callable mapping unconstrained site values to the potential energy. + """ + if self._potential_fn_gen is None: + if self._potential_fn is not None: + return self._potential_fn + raise RuntimeError( + "`get_potential_fn` requires the kernel to be initialized; run `init` first." + ) + model_kwargs = {} if model_kwargs is None else model_kwargs + return self._potential_fn_gen(*model_args, **model_kwargs) + + def get_constrain_fn( + self, + model_args: ModelArgs = (), + model_kwargs: ModelKwargs | None = None, + *, + return_deterministic: bool = False, + ) -> ConstrainFn: + """ + Return a function mapping unconstrained sample values to constrained values. When + `return_deterministic=False` and the model has no value-dependent supports, this is a + transform-only function (:func:`~numpyro.infer.util.transform_fn`) that never runs the + model; otherwise it replays the model (:func:`~numpyro.infer.util.constrain_fn`). + Composite kernels use the transform-only form to condition sibling blocks. Requires + :meth:`init` to have run. + + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model. + :param bool return_deterministic: whether to also return `deterministic` sites. + :return: a callable mapping unconstrained site values to constrained site values. + """ + if self._inv_transforms is None: + if self._model is None: + return identity + raise RuntimeError( + "`get_constrain_fn` requires the kernel to be initialized; run `init` first." + ) + if return_deterministic or self._dynamic_support: + model_kwargs = {} if model_kwargs is None else model_kwargs + assert self._postprocess_fn is not None + return self._postprocess_fn(*model_args, **model_kwargs) + return partial(transform_fn, self._inv_transforms) + + def refresh( + self, + state: HMCState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> HMCState: + """ + Recompute `potential_energy` and `z_grad` at `state.z` with the potential for the + given arguments, honoring `forward_mode_differentiation`. `energy` is left as is: + :meth:`sample` recomputes it from the potential and a fresh momentum. + + :param HMCState state: the current state. + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model. + :return: the state with refreshed `potential_energy` and `z_grad`. + """ + pe_fn = self.get_potential_fn(model_args, model_kwargs) + pe, z_grad = _value_and_grad(pe_fn, state.z, self._forward_mode_differentiation) + return state._replace(potential_energy=pe, z_grad=z_grad) + + def wrap_model(self, wrapper: Callable) -> "HMC": + """ + Return a shallow copy of this kernel bound to `wrapper(self.model)`. The copy rebuilds + its closures against the wrapped model on its next :meth:`init`. + + :param wrapper: callable mapping a model to a model with the same call signature. + :return: a new kernel bound to the wrapped model. + """ + if self._model is None: + raise ValueError( + "`wrap_model` is not supported for kernels built from a potential function." + ) + kernel = copy.copy(self) + kernel._model = wrapper(self._model) + kernel._init_fn = None + kernel._sample_fn = None + kernel._potential_fn_gen = None + kernel._postprocess_fn = None + kernel._inv_transforms = None + kernel._dynamic_support = None + return kernel + def init( self, rng_key, num_warmup, init_params=None, model_args=(), model_kwargs={} ): @@ -768,7 +886,9 @@ def init( dense_mass = [tuple(sorted(z))] if dense_mass else [] assert isinstance(dense_mass, list) - hmc_init_fn = lambda init_params, rng_key: self._init_fn( # noqa: E731 + init_fn = self._init_fn + assert init_fn is not None + hmc_init_fn = lambda init_params, rng_key: init_fn( # noqa: E731 init_params, num_warmup=num_warmup, step_size=self._step_size, @@ -794,14 +914,15 @@ def init( # nonlocal variables: momentum_generator, wa_update, trajectory_len, max_treedepth, # wa_steps because those variables do not depend on traced args: init_params, rng_key. init_state = vmap(hmc_init_fn)(init_params, rng_key) + assert self._sample_fn is not None sample_fn = vmap(self._sample_fn, in_axes=(0, None, None)) self._sample_fn = sample_fn return init_state - def postprocess_fn(self, args, kwargs): + def postprocess_fn(self, model_args, model_kwargs): if self._postprocess_fn is None: return identity - return self._postprocess_fn(*args, **kwargs) + return self._postprocess_fn(*model_args, **model_kwargs) def sample(self, state, model_args, model_kwargs): """ @@ -813,6 +934,7 @@ def sample(self, state, model_args, model_kwargs): :param model_kwargs: Keyword arguments provided to the model. :return: Next `state` after running HMC. """ + assert self._sample_fn is not None, "`init` must be called before `sample`." return self._sample_fn(state, model_args, model_kwargs) def __getstate__(self): diff --git a/numpyro/infer/hmc_gibbs.py b/numpyro/infer/hmc_gibbs.py index ac5032904..cd5bdcdf2 100644 --- a/numpyro/infer/hmc_gibbs.py +++ b/numpyro/infer/hmc_gibbs.py @@ -2,46 +2,72 @@ # SPDX-License-Identifier: Apache-2.0 from collections import namedtuple +from collections.abc import Callable, Sequence import copy from functools import partial +from typing import Any, TypeAlias -import numpy as np - -from jax import grad, jacfwd, random, value_and_grad -from jax.flatten_util import ravel_pytree +import jax +from jax import random import jax.numpy as jnp -from jax.scipy.special import expit import numpyro +from numpyro._typing import ( + ConstrainFn, + ModelArgs, + ModelKwargs, + ModelT, + PyTree, + SiteValues, +) from numpyro.contrib.ecs_proxies import block_update, perturbed_method, taylor_proxy -from numpyro.handlers import condition, seed, substitute, trace -from numpyro.infer.hmc import HMC -from numpyro.infer.initialization import init_to_sample +from numpyro.infer.gibbs import ( + GIBBS_SITES_KWARG, + CustomGibbs, + DiscreteGibbs, + Gibbs, + GibbsState, + GibbsUpdateFn, + ModelWrapper, + conditioned, + discrete_latent_sites, + prototype_trace, + subsample_plate_sizes, + with_conditioning, +) +from numpyro.infer.hmc import HMC, HMCState from numpyro.infer.mcmc import MCMCKernel -from numpyro.infer.util import _unconstrain_reparam -from numpyro.util import cond, fori_loop, identity +from numpyro.infer.util import _unconstrain_params +from numpyro.util import cond, identity -HMCGibbsState = namedtuple("HMCGibbsState", "z, hmc_state, rng_key") -""" - - **z** - a dict of the current latent values (both HMC and Gibbs sites) - - **hmc_state** - current :data:`~numpyro.infer.hmc.HMCState` - - **rng_key** - random key for the current step -""" +class HMCGibbsState(GibbsState): + """ + :class:`~numpyro.infer.gibbs.GibbsState` of :class:`HMCGibbs` and :class:`DiscreteHMCGibbs`, + constructed by their `init` method (not positionally). + + - **z** - a dict of the current latent values (both HMC and Gibbs sites) + - **block_states** - the states of the Gibbs block and of the HMC block + - **rng_key** - random key for the current step + - **hmc_state** - property returning the current :data:`~numpyro.infer.hmc.HMCState` + (the last block state), so that `extra_fields=["hmc_state.potential_energy"]` works + """ -def _wrap_model(model, *args, **kwargs): - gibbs_values = kwargs.pop("_gibbs_sites", {}) - with condition(data=gibbs_values), substitute(data=gibbs_values): - return model(*args, **kwargs) + __slots__ = () + + @property + def hmc_state(self) -> HMCState: + return self.block_states[-1] -class HMCGibbs(MCMCKernel): +class HMCGibbs(Gibbs): """ [EXPERIMENTAL INTERFACE] HMC-within-Gibbs. This inference algorithm allows the user to combine general purpose gradient-based inference (HMC or NUTS) with custom - Gibbs samplers. + Gibbs samplers. It is equivalent to + ``Gibbs([(CustomGibbs(gibbs_fn), gibbs_sites), (inner_kernel, None)])``. Note that it is the user's responsibility to provide a correct implementation of `gibbs_fn` that samples from the corresponding posterior conditional. @@ -82,270 +108,29 @@ class HMCGibbs(MCMCKernel): """ - sample_field = "z" + _state_cls = HMCGibbsState - def __init__(self, inner_kernel, gibbs_fn, gibbs_sites): + def __init__( + self, + inner_kernel: HMC, + gibbs_fn: GibbsUpdateFn, + gibbs_sites: Sequence[str], + ) -> None: if not isinstance(inner_kernel, HMC): raise ValueError("inner_kernel must be an HMC or NUTS sampler.") if not callable(gibbs_fn): raise ValueError("gibbs_fn must be a callable") - assert inner_kernel.model is not None, ( - "HMCGibbs does not support models specified via a potential function." - ) - - self.inner_kernel = copy.copy(inner_kernel) - self.inner_kernel._model = partial(_wrap_model, inner_kernel.model) - self._gibbs_sites = gibbs_sites - self._gibbs_fn = gibbs_fn - self._prototype_trace = None - - @property - def model(self): - return self.inner_kernel._model - - def get_diagnostics_str(self, state): - state = state.hmc_state - return "{} steps of size {:.2e}. acc. prob={:.2f}".format( - state.num_steps, state.adapt_state.step_size, state.mean_accept_prob - ) - - def postprocess_fn(self, args, kwargs): - def fn(z): - model_kwargs = {} if kwargs is None else kwargs.copy() - hmc_sites = {k: v for k, v in z.items() if k not in self._gibbs_sites} - gibbs_sites = {k: v for k, v in z.items() if k in self._gibbs_sites} - model_kwargs["_gibbs_sites"] = gibbs_sites - hmc_sites = self.inner_kernel.postprocess_fn(args, model_kwargs)(hmc_sites) - return {**gibbs_sites, **hmc_sites} - - return fn - - def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs): - model_kwargs = {} if model_kwargs is None else model_kwargs.copy() - if self._prototype_trace is None: - rng_key, key_u = random.split(rng_key) - # We use init strategy to get around ImproperUniform which does not have - # sample method. - self._prototype_trace = trace( - substitute(seed(self.model, key_u), substitute_fn=init_to_sample) - ).get_trace(*model_args, **model_kwargs) - - rng_key, key_z = random.split(rng_key) - - gibbs_sites = {} - - for name, site in self._prototype_trace.items(): - if init_params and (name in init_params) and (name in self._gibbs_sites): - gibbs_sites[name] = init_params.pop(name) - - elif name in self._gibbs_sites: - gibbs_sites[name] = site["value"] - - model_kwargs["_gibbs_sites"] = gibbs_sites - hmc_state = self.inner_kernel.init( - key_z, num_warmup, init_params, model_args, model_kwargs - ) - - z = {**gibbs_sites, **hmc_state.z} - - return HMCGibbsState(z, hmc_state, rng_key) - - def sample(self, state, model_args, model_kwargs): - model_kwargs = {} if model_kwargs is None else model_kwargs - rng_key, rng_gibbs = random.split(state.rng_key) - - def potential_fn(z_gibbs, z_hmc): - return self.inner_kernel._potential_fn_gen( - *model_args, _gibbs_sites=z_gibbs, **model_kwargs - )(z_hmc) - - z_gibbs = {k: v for k, v in state.z.items() if k not in state.hmc_state.z} - z_hmc = {k: v for k, v in state.z.items() if k in state.hmc_state.z} - model_kwargs_ = model_kwargs.copy() - model_kwargs_["_gibbs_sites"] = z_gibbs - z_hmc = self.inner_kernel.postprocess_fn(model_args, model_kwargs_)(z_hmc) - - z_gibbs = self._gibbs_fn( - rng_key=rng_gibbs, gibbs_sites=z_gibbs, hmc_sites=z_hmc - ) - - if self.inner_kernel._forward_mode_differentiation: - pe = potential_fn(z_gibbs, state.hmc_state.z) - z_grad = jacfwd(partial(potential_fn, z_gibbs))(state.hmc_state.z) - else: - pe, z_grad = value_and_grad(partial(potential_fn, z_gibbs))( - state.hmc_state.z - ) - hmc_state = state.hmc_state._replace(z_grad=z_grad, potential_energy=pe) - - model_kwargs_["_gibbs_sites"] = z_gibbs - hmc_state = self.inner_kernel.sample(hmc_state, model_args, model_kwargs_) - - z = {**z_gibbs, **hmc_state.z} - - return HMCGibbsState(z, hmc_state, rng_key) - - def __getstate__(self): - state = self.__dict__.copy() - state["_prototype_trace"] = None - return state - - -def _discrete_gibbs_proposal_body_fn( - z_init_flat, unravel_fn, pe_init, potential_fn, idx, i, val -): - rng_key, z, pe, log_weight_sum = val - rng_key, rng_transition = random.split(rng_key) - proposal = jnp.where(i >= z_init_flat[idx], i + 1, i) - z_new_flat = z_init_flat.at[idx].set(proposal) - z_new = unravel_fn(z_new_flat) - pe_new = potential_fn(z_new) - log_weight_new = pe_init - pe_new - # Handles the NaN case... - log_weight_new = jnp.where(jnp.isfinite(log_weight_new), log_weight_new, -jnp.inf) - # transition_prob = e^weight_new / (e^weight_logsumexp + e^weight_new) - transition_prob = expit(log_weight_new - log_weight_sum) - z, pe = cond( - random.bernoulli(rng_transition, transition_prob), - (z_new, pe_new), - identity, - (z, pe), - identity, - ) - log_weight_sum = jnp.logaddexp(log_weight_new, log_weight_sum) - return rng_key, z, pe, log_weight_sum - - -def _discrete_gibbs_proposal(rng_key, z_discrete, pe, potential_fn, idx, support_size): - # idx: current index of `z_discrete_flat` to update - # support_size: support size of z_discrete at the index idx - - z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) - # Here we loop over the support of z_flat[idx] to get z_new - # Note: we can't vmap potential_fn over all proposals and sample from the conditional - # categorical distribution because support_size is a traced value, i.e. its value - # might change across different discrete variables; - # so here we will loop over all proposals and use an online scheme to sample from - # the conditional categorical distribution - body_fn = partial( - _discrete_gibbs_proposal_body_fn, - z_discrete_flat, - unravel_fn, - pe, - potential_fn, - idx, - ) - init_val = (rng_key, z_discrete, pe, jnp.array(0.0)) - rng_key, z_new, pe_new, _ = fori_loop(0, support_size - 1, body_fn, init_val) - log_accept_ratio = jnp.array(0.0) - return rng_key, z_new, pe_new, log_accept_ratio - - -def _discrete_modified_gibbs_proposal( - rng_key, z_discrete, pe, potential_fn, idx, support_size, stay_prob=0.0 -): - assert isinstance(stay_prob, float) and stay_prob >= 0.0 and stay_prob < 1 - z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) - body_fn = partial( - _discrete_gibbs_proposal_body_fn, - z_discrete_flat, - unravel_fn, - pe, - potential_fn, - idx, - ) - # like gibbs_step but here, weight of the current value is 0 - init_val = (rng_key, z_discrete, pe, jnp.array(-jnp.inf)) - rng_key, z_new, pe_new, log_weight_sum = fori_loop( - 0, support_size - 1, body_fn, init_val - ) - rng_key, rng_stay = random.split(rng_key) - z_new, pe_new = cond( - random.bernoulli(rng_stay, stay_prob), - (z_discrete, pe), - identity, - (z_new, pe_new), - identity, - ) - # here we calculate the MH correction: (1 - P(z)) / (1 - P(z_new)) - # where 1 - P(z) ~ weight_sum - # and 1 - P(z_new) ~ 1 + weight_sum - z_new_weight - log_accept_ratio = log_weight_sum - jnp.log( - jnp.exp(log_weight_sum) - jnp.expm1(pe - pe_new) - ) - return rng_key, z_new, pe_new, log_accept_ratio - - -def _discrete_rw_proposal(rng_key, z_discrete, pe, potential_fn, idx, support_size): - rng_key, rng_proposal = random.split(rng_key, 2) - z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) - - proposal = random.randint(rng_proposal, (), minval=0, maxval=support_size) - z_new_flat = z_discrete_flat.at[idx].set(proposal) - z_new = unravel_fn(z_new_flat) - pe_new = potential_fn(z_new) - log_accept_ratio = pe - pe_new - return rng_key, z_new, pe_new, log_accept_ratio - - -def _discrete_modified_rw_proposal( - rng_key, z_discrete, pe, potential_fn, idx, support_size, stay_prob=0.0 -): - assert isinstance(stay_prob, float) and stay_prob >= 0.0 and stay_prob < 1 - rng_key, rng_proposal, rng_stay = random.split(rng_key, 3) - z_discrete_flat, unravel_fn = ravel_pytree(z_discrete) - - i = random.randint(rng_proposal, (), minval=0, maxval=support_size - 1) - proposal = jnp.where(i >= z_discrete_flat[idx], i + 1, i) - proposal = jnp.where(random.bernoulli(rng_stay, stay_prob), idx, proposal) - z_new_flat = z_discrete_flat.at[idx].set(proposal) - z_new = unravel_fn(z_new_flat) - pe_new = potential_fn(z_new) - log_accept_ratio = pe - pe_new - return rng_key, z_new, pe_new, log_accept_ratio - - -def _discrete_gibbs_fn(potential_fn, support_sizes, proposal_fn): - def gibbs_fn(rng_key, gibbs_sites, hmc_sites, pe): - # get support_sizes of gibbs_sites - support_sizes_flat, _ = ravel_pytree({k: support_sizes[k] for k in gibbs_sites}) - num_discretes = support_sizes_flat.shape[0] - - rng_key, rng_permute = random.split(rng_key) - idxs = random.permutation(rng_key, jnp.arange(num_discretes)) - - def body_fn(i, val): - idx = idxs[i] - support_size = support_sizes_flat[idx] - rng_key, z, pe = val - rng_key, z_new, pe_new, log_accept_ratio = proposal_fn( - rng_key, - z, - pe, - potential_fn=partial(potential_fn, z_hmc=hmc_sites), - idx=idx, - support_size=support_size, + if inner_kernel.model is None: + raise ValueError( + "HMCGibbs does not support models specified via a potential function." ) - rng_key, rng_accept = random.split(rng_key) - # u ~ Uniform(0, 1), u < accept_ratio => -log(u) > -log_accept_ratio - # and -log(u) ~ exponential(1) - z, pe = cond( - random.exponential(rng_accept) > -log_accept_ratio, - (z_new, pe_new), - identity, - (z, pe), - identity, - ) - return rng_key, z, pe - - init_val = (rng_key, gibbs_sites, pe) - _, gibbs_sites, pe = fori_loop(0, num_discretes, body_fn, init_val) - return gibbs_sites, pe - - return gibbs_fn + super().__init__([(CustomGibbs(gibbs_fn), gibbs_sites), (inner_kernel, None)]) + self.inner_kernel = self._kernels[1] + self._gibbs_fn = gibbs_fn + self._gibbs_sites = tuple(gibbs_sites) -class DiscreteHMCGibbs(HMCGibbs): +class DiscreteHMCGibbs(Gibbs): """ [EXPERIMENTAL INTERFACE] @@ -392,123 +177,87 @@ class DiscreteHMCGibbs(HMCGibbs): >>> mcmc.run(random.key(0), probs, locs) >>> mcmc.print_summary() # doctest: +SKIP >>> samples = mcmc.get_samples()["x"] - >>> assert abs(jnp.mean(samples) - 1.3) < 0.1 + >>> assert abs(jnp.mean(samples) - 1.3) < 0.2 >>> assert abs(jnp.var(samples) - 4.36) < 0.5 """ - def __init__(self, inner_kernel, *, random_walk=False, modified=False): - super().__init__(inner_kernel, identity, None) - self._random_walk = random_walk - self._modified = modified - if random_walk: - if modified: - self._discrete_proposal_fn = partial( - _discrete_modified_rw_proposal, stay_prob=0.0 - ) - else: - self._discrete_proposal_fn = _discrete_rw_proposal - else: - if modified: - self._discrete_proposal_fn = partial( - _discrete_modified_gibbs_proposal, stay_prob=0.0 - ) - else: - self._discrete_proposal_fn = _discrete_gibbs_proposal + _state_cls = HMCGibbsState - def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs): - model_kwargs = {} if model_kwargs is None else model_kwargs.copy() - rng_key, key_u = random.split(rng_key) - # We use init strategy to get around ImproperUniform which does not have - # sample method. - self._prototype_trace = trace( - substitute(seed(self.model, key_u), substitute_fn=init_to_sample) - ).get_trace(*model_args, **model_kwargs) - - self._support_sizes = { - name: np.broadcast_to( - site["fn"].enumerate_support(False).shape[0], jnp.shape(site["value"]) + def __init__( + self, + inner_kernel: HMC, + *, + random_walk: bool = False, + modified: bool = False, + ) -> None: + if not isinstance(inner_kernel, HMC): + raise ValueError("inner_kernel must be an HMC or NUTS sampler.") + if inner_kernel.model is None: + raise ValueError( + "DiscreteHMCGibbs does not support models specified via a potential function." ) - for name, site in self._prototype_trace.items() - if site["type"] == "sample" - and site["fn"].has_enumerate_support - and not site["is_observed"] - } - self._gibbs_sites = [ - name - for name, site in self._prototype_trace.items() - if site["type"] == "sample" - and site["fn"].has_enumerate_support - and not site["is_observed"] - and site["infer"].get("enumerate", "") != "parallel" - ] - assert self._gibbs_sites, ( - "Cannot detect any discrete latent variables in the model." + discrete_kernel = DiscreteGibbs( + inner_kernel.model, random_walk=random_walk, modified=modified ) - return super().init(rng_key, num_warmup, init_params, model_args, model_kwargs) - - def sample(self, state, model_args, model_kwargs): - model_kwargs = {} if model_kwargs is None else model_kwargs - rng_key, rng_gibbs = random.split(state.rng_key) - - def potential_fn(z_gibbs, z_hmc): - return self.inner_kernel._potential_fn_gen( - *model_args, _gibbs_sites=z_gibbs, **model_kwargs - )(z_hmc) - - z_gibbs = {k: v for k, v in state.z.items() if k not in state.hmc_state.z} - z_hmc = {k: v for k, v in state.z.items() if k in state.hmc_state.z} - model_kwargs_ = model_kwargs.copy() - model_kwargs_["_gibbs_sites"] = z_gibbs - - # different from the implementation in HMCGibbs.sample, we feed the current potential energy - # and get new potential energy from gibbs_fn - gibbs_fn = _discrete_gibbs_fn( - potential_fn, self._support_sizes, self._discrete_proposal_fn - ) - z_gibbs, pe = gibbs_fn( - rng_key=rng_gibbs, - gibbs_sites=z_gibbs, - hmc_sites=z_hmc, - pe=state.hmc_state.potential_energy, + super().__init__( + [(discrete_kernel, discrete_latent_sites), (inner_kernel, None)] ) + self.inner_kernel = self._kernels[1] + self._random_walk = random_walk + self._modified = modified - if self.inner_kernel._forward_mode_differentiation: - z_grad = jacfwd(partial(potential_fn, z_gibbs))(state.hmc_state.z) - else: - z_grad = grad(partial(potential_fn, z_gibbs))(state.hmc_state.z) - hmc_state = state.hmc_state._replace(z_grad=z_grad, potential_energy=pe) - model_kwargs_["_gibbs_sites"] = z_gibbs - hmc_state = self.inner_kernel.sample(hmc_state, model_args, model_kwargs_) +HMCECSState = namedtuple( + "HMCECSState", "z, hmc_state, rng_key, gibbs_state, accept_prob" +) - z = {**z_gibbs, **hmc_state.z} +LikelihoodEstimator: TypeAlias = Callable[ + [dict[str, tuple], SiteValues, PyTree], jax.Array +] +""" +`(likelihoods, unconstrained_params, gibbs_state) -> log-likelihood estimate`; see +`perturbed_method` in :mod:`numpyro.contrib.ecs_proxies`. +""" - return HMCGibbsState(z, hmc_state, rng_key) +ProxyConstructor: TypeAlias = Callable[ + ..., tuple[Callable[..., Any], Callable[..., Any], Callable[..., Any]] +] +""" +`(prototype_trace, subsample_plate_sizes, model, model_args, model_kwargs, num_blocks) -> +(proxy_fn, gibbs_init, gibbs_update)`; see :func:`~numpyro.contrib.ecs_proxies.taylor_proxy`. +""" -HMCECSState = namedtuple( - "HMCECSState", "z, hmc_state, rng_key, gibbs_state, accept_prob" -) +def _ecs_model(model, estimator, *args, **kwargs): + """ + Model wrapper installed once by :class:`HMCECS`: pops `_gibbs_state` from the keyword + arguments, hands it to `estimator`, and runs `model` under `estimator`. When the estimator + has no `method` (no proxy), the model runs with its plain subsampled likelihood. + """ + gibbs_state = kwargs.pop("_gibbs_state", ()) + if estimator.method is None: + return model(*args, **kwargs) + estimator.gibbs_state = gibbs_state + with estimator: + return model(*args, **kwargs) -def _wrap_gibbs_state(model, *args, **kwargs): - # this is to let estimate_likelihood handler knows what is the current gibbs_state - msg = {"type": "_gibbs_state", "value": kwargs.pop("_gibbs_state", ())} - numpyro.primitives.apply_stack(msg) - return model(*args, **kwargs) +def _wrap_ecs(model, estimator): + return partial(_ecs_model, conditioned(model), estimator) -class HMCECS(HMCGibbs): +class HMCECS(MCMCKernel): """ [EXPERIMENTAL INTERFACE] HMC with Energy Conserving Subsampling. - A subclass of :class:`HMCGibbs` for performing HMC-within-Gibbs for models with subsample - statements using the :class:`~numpyro.plate` primitive. This implements Algorithm 1 - of reference [1] but uses a naive estimation (without control variates) of log likelihood, - hence might incur a high variance. + A wrapper around an HMC kernel for performing HMC-within-Gibbs for models with subsample + statements using the :class:`~numpyro.plate` primitive: it changes the target of the + inner kernel (likelihood estimator) and performs the pseudo-marginal Metropolis update of + the subsample indices. This implements Algorithm 1 of reference [1] but uses a naive + estimation (without control variates) of log likelihood, hence might incur a high variance. The function can divide subsample indices into blocks and update only one block at each MCMC step to improve the acceptance rate of proposed subsamples as detailed in [3]. @@ -552,149 +301,275 @@ class HMCECS(HMCGibbs): >>> mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000) >>> mcmc.run(random.key(0), data) >>> samples = mcmc.get_samples()["x"] - >>> assert abs(jnp.mean(samples) - 1.) < 0.1 + >>> assert abs(jnp.mean(samples) - 1.) < 0.2 """ - def __init__(self, inner_kernel, *, num_blocks=1, proxy=None): - super().__init__(inner_kernel, identity, None) + sample_field: str = "z" - self.inner_kernel._model = partial(_wrap_gibbs_state, self.inner_kernel._model) + def __init__( + self, + inner_kernel: HMC, + *, + num_blocks: int = 1, + proxy: ProxyConstructor | None = None, + ) -> None: + if not isinstance(inner_kernel, HMC): + raise ValueError("inner_kernel must be an HMC or NUTS sampler.") + if inner_kernel.model is None: + raise ValueError( + "HMCECS does not support models specified via a potential function." + ) + self._model = inner_kernel.model + self._estimator = estimate_likelihood() + # wrap once, at construction: `init` only binds the estimator's method and state + self.inner_kernel = inner_kernel.wrap_model( + partial(_wrap_ecs, estimator=self._estimator) + ) self._num_blocks = num_blocks self._proxy = proxy + # static metadata resolved at `init` + self._subsample_plate_sizes: dict[str, tuple[int, int]] | None = None + self._gibbs_sites: tuple[str, ...] = () + self._gibbs_update = None + self._sample_fn = None + + @property + def model(self) -> ModelT: + return self._model + + def get_diagnostics_str(self, state: HMCECSState) -> str: + return self.inner_kernel.get_diagnostics_str(state.hmc_state) + + def _inner_kwargs( + self, model_kwargs: ModelKwargs | None, z_gibbs: SiteValues, gibbs_state: PyTree + ) -> ModelKwargs: + model_kwargs = with_conditioning(model_kwargs, z_gibbs) + model_kwargs["_gibbs_state"] = gibbs_state + return model_kwargs + + def _split(self, z: SiteValues) -> tuple[SiteValues, SiteValues]: + z_gibbs = {k: v for k, v in z.items() if k in self._gibbs_sites} + z_hmc = {k: v for k, v in z.items() if k not in self._gibbs_sites} + return z_gibbs, z_hmc + + def postprocess_fn( + self, model_args: ModelArgs, model_kwargs: ModelKwargs | None + ) -> ConstrainFn: + """Inner postprocess on the HMC sites; subsample indices are dropped.""" + + def fn(z: SiteValues) -> SiteValues: + z_gibbs, z_hmc = self._split(z) + return self.inner_kernel.postprocess_fn( + model_args, with_conditioning(model_kwargs, z_gibbs) + )(z_hmc) - def postprocess_fn(self, args, kwargs): - def fn(z): - model_kwargs = {} if kwargs is None else kwargs.copy() - hmc_sites = {k: v for k, v in z.items() if k not in self._gibbs_sites} - gibbs_sites = {k: v for k, v in z.items() if k in self._gibbs_sites} - model_kwargs["_gibbs_sites"] = gibbs_sites - hmc_sites = self.inner_kernel.postprocess_fn(args, model_kwargs)(hmc_sites) - return hmc_sites + return fn + + def get_constrain_fn( + self, model_args: ModelArgs, model_kwargs: ModelKwargs | None + ) -> ConstrainFn: + """Inner constrain function on the HMC sites; subsample indices are dropped.""" + + def fn(z: SiteValues) -> SiteValues: + z_gibbs, z_hmc = self._split(z) + return self.inner_kernel.get_constrain_fn( + model_args, with_conditioning(model_kwargs, z_gibbs) + )(z_hmc) return fn - def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs): - model_kwargs = {} if model_kwargs is None else model_kwargs.copy() + def init( + self, + rng_key: jax.Array, + num_warmup: int, + init_params: SiteValues | None, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> HMCECSState: + model_kwargs = {} if model_kwargs is None else dict(model_kwargs) rng_key, key_u = random.split(rng_key) - # We use init strategy to get around ImproperUniform which does not have - # sample method. - self._prototype_trace = trace( - substitute(seed(self.model, key_u), substitute_fn=init_to_sample) - ).get_trace(*model_args, **model_kwargs) - self._subsample_plate_sizes = { - name: site["args"] - for name, site in self._prototype_trace.items() - if site["type"] == "plate" - and (site["args"][1] is not None) - and site["args"][0] > site["args"][1] - } # i.e. size > subsample_size - self._gibbs_sites = list(self._subsample_plate_sizes.keys()) - assert self._gibbs_sites, "Cannot detect any subsample statements in the model." + model = conditioned(self._model) + model_trace = prototype_trace(model, key_u, model_args, model_kwargs) + self._subsample_plate_sizes = subsample_plate_sizes(model_trace) + self._gibbs_sites = tuple(self._subsample_plate_sizes) + if not self._gibbs_sites: + raise ValueError("Cannot detect any subsample statements in the model.") + for name in model_kwargs.get(GIBBS_SITES_KWARG, {}): + site = model_trace.get(name) + if site is not None and any( + frame.name in self._subsample_plate_sizes + for frame in site["cond_indep_stack"] + ): + raise ValueError( + f"Site '{name}' is conditioned by an enclosing kernel but lies inside " + "a subsample plate; HMCECS cannot estimate its likelihood." + ) if self._proxy is not None: if any( - { - name - for name, site in self._prototype_trace.items() - if site["type"] == "sample" - and (not site["is_observed"]) - and site["fn"].support.is_discrete - } + site["type"] == "sample" + and (not site["is_observed"]) + and site["fn"].support.is_discrete + for site in model_trace.values() ): raise RuntimeError( "Currently, the proxy does not support models with " "discrete latent sites." ) proxy_fn, gibbs_init, self._gibbs_update = self._proxy( - self._prototype_trace, + model_trace, self._subsample_plate_sizes, - self.model, + model, model_args, model_kwargs.copy(), num_blocks=self._num_blocks, ) - method = perturbed_method(self._subsample_plate_sizes, proxy_fn) - self.inner_kernel._model = estimate_likelihood( - self.inner_kernel._model, method + self._estimator.method = perturbed_method( + self._subsample_plate_sizes, proxy_fn ) - - z_gibbs = { - name: site["value"] - for name, site in self._prototype_trace.items() - if name in self._gibbs_sites - } - rng_key, rng_state = random.split(rng_key) - gibbs_state = gibbs_init(rng_state, z_gibbs) else: + self._estimator.method = None self._gibbs_update = partial( block_update, self._subsample_plate_sizes, self._num_blocks ) - gibbs_state = () - model_kwargs["_gibbs_state"] = gibbs_state - state = super().init(rng_key, num_warmup, init_params, model_args, model_kwargs) - return HMCECSState( - state.z, state.hmc_state, state.rng_key, gibbs_state, jnp.zeros(()) - ) - - def sample(self, state, model_args, model_kwargs): - model_kwargs = {} if model_kwargs is None else model_kwargs.copy() - rng_key, rng_gibbs = random.split(state.rng_key) + init_params = None if init_params is None else dict(init_params) + z_gibbs = {} + for name in self._gibbs_sites: + if init_params and name in init_params: + z_gibbs[name] = init_params.pop(name) + else: + z_gibbs[name] = model_trace[name]["value"] - def potential_fn(z_gibbs, gibbs_state, z_hmc): - return self.inner_kernel._potential_fn_gen( - *model_args, - _gibbs_sites=z_gibbs, - _gibbs_state=gibbs_state, - **model_kwargs, - )(z_hmc) + if self._proxy is not None: + rng_key, rng_state = random.split(rng_key) + gibbs_state = gibbs_init(rng_state, z_gibbs) + else: + gibbs_state = () - z_gibbs = {k: v for k, v in state.z.items() if k not in state.hmc_state.z} + rng_key, key_z = random.split(rng_key) + hmc_state = self.inner_kernel.init( + key_z, + num_warmup, + init_params or None, + model_args, + self._inner_kwargs(model_kwargs, z_gibbs, gibbs_state), + ) + z = {**z_gibbs, **hmc_state.z} + self._sample_fn = self._sample_one + return HMCECSState(z, hmc_state, rng_key, gibbs_state, jnp.zeros(())) + + def _sample_one( + self, + state: HMCECSState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> HMCECSState: + rng_key, rng_gibbs, rng_accept = random.split(state.rng_key, 3) + + z_gibbs, _ = self._split(state.z) + assert self._gibbs_update is not None, "`init` must be called before `sample`." z_gibbs_new, gibbs_state_new = self._gibbs_update( - rng_key, z_gibbs, state.gibbs_state + rng_gibbs, z_gibbs, state.gibbs_state ) # given a fixed hmc_sites, pe_new - pe_curr = loglik_new - loglik_curr pe = state.hmc_state.potential_energy - pe_new = potential_fn(z_gibbs_new, gibbs_state_new, state.hmc_state.z) + pe_new = self.inner_kernel.get_potential_fn( + model_args, self._inner_kwargs(model_kwargs, z_gibbs_new, gibbs_state_new) + )(state.hmc_state.z) accept_prob = jnp.clip(jnp.exp(pe - pe_new), None, 1.0) - transition = random.bernoulli(rng_key, accept_prob) - grad_ = jacfwd if self.inner_kernel._forward_mode_differentiation else grad + transition = random.bernoulli(rng_accept, accept_prob) + + def accept(vals): + z_gibbs_new, gibbs_state_new, _ = vals + refreshed = self.inner_kernel.refresh( + state.hmc_state, + model_args, + self._inner_kwargs(model_kwargs, z_gibbs_new, gibbs_state_new), + ) + return ( + z_gibbs_new, + gibbs_state_new, + refreshed.potential_energy, + refreshed.z_grad, + ) + z_gibbs, gibbs_state, pe, z_grad = cond( transition, (z_gibbs_new, gibbs_state_new, pe_new), - lambda vals: ( - vals - + (grad_(partial(potential_fn, vals[0], vals[1]))(state.hmc_state.z),) - ), + accept, (z_gibbs, state.gibbs_state, pe, state.hmc_state.z_grad), identity, ) hmc_state = state.hmc_state._replace(z_grad=z_grad, potential_energy=pe) - - model_kwargs["_gibbs_sites"] = z_gibbs - model_kwargs["_gibbs_state"] = gibbs_state - hmc_state = self.inner_kernel.sample(hmc_state, model_args, model_kwargs) + hmc_state = self.inner_kernel.sample( + hmc_state, + model_args, + self._inner_kwargs(model_kwargs, z_gibbs, gibbs_state), + ) z = {**z_gibbs, **hmc_state.z} return HMCECSState(z, hmc_state, rng_key, gibbs_state, accept_prob) + def sample( + self, + state: HMCECSState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> HMCECSState: + assert self._sample_fn is not None, "`init` must be called before `sample`." + return self._sample_fn(state, model_args, model_kwargs) + + def refresh( + self, + state: HMCECSState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> HMCECSState: + """Delegates to the inner kernel with the current subsample indices and proxy state.""" + z_gibbs, _ = self._split(state.z) + hmc_state = self.inner_kernel.refresh( + state.hmc_state, + model_args, + self._inner_kwargs(model_kwargs, z_gibbs, state.gibbs_state), + ) + return state._replace(hmc_state=hmc_state) + + def wrap_model(self, wrapper: ModelWrapper) -> "HMCECS": + kernel = copy.copy(self) + kernel._model = wrapper(self._model) + kernel.inner_kernel = self.inner_kernel.wrap_model(wrapper) + kernel._sample_fn = None + return kernel + @staticmethod - def taylor_proxy(reference_params, degree=2): + def taylor_proxy(reference_params: SiteValues, degree: int = 2) -> ProxyConstructor: """ This is just a convenient static method which calls :func:`~numpyro.contrib.ecs_proxies.taylor_proxy`. """ return taylor_proxy(reference_params, degree) + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state["_sample_fn"] = None + return state + class estimate_likelihood(numpyro.primitives.Messenger): - def __init__(self, fn=None, method=None): - # estimate_likelihood: accept likelihood tuple (fn, value, subsample_name, subsample_dim) - # and current unconstrained params - # and returns log of the bias-corrected likelihood - assert method is not None + """ + Handler that replaces the subsampled likelihood of a model by a bias-corrected estimate. + `method` accepts the likelihood tuples `(fn, value, subsample_name, subsample_dim)`, the + current unconstrained parameters and the proxy state (`gibbs_state`) and returns the log + of the estimated likelihood. Both `method` and `gibbs_state` can be set after + construction; the handler is inert while `method` is `None`. + """ + + def __init__( + self, fn: ModelT | None = None, method: LikelihoodEstimator | None = None + ) -> None: super().__init__(fn) self.method = method self.params = None @@ -703,16 +578,13 @@ def __init__(self, fn=None, method=None): self.gibbs_state = None def __enter__(self): - for handler in numpyro.primitives._PYRO_STACK[::-1]: - # the potential_fn in HMC makes the PYRO_STACK nested like trace(...); so we can extract the - # unconstrained_params from the _unconstrain_reparam substitute_fn - if ( - isinstance(handler, substitute) - and isinstance(handler.substitute_fn, partial) - and handler.substitute_fn.func is _unconstrain_reparam - ): - self.params = handler.substitute_fn.args[0] - break + if self.method is not None: + for handler in numpyro.primitives._PYRO_STACK[::-1]: + # the potential_fn in HMC makes the PYRO_STACK nested like trace(...); so we + # can extract the unconstrained params from the `_unconstrain_params` handler + if isinstance(handler, _unconstrain_params): + self.params = handler.params + break return super().__enter__() def __exit__(self, exc_type, exc_value, traceback): @@ -725,6 +597,7 @@ def __exit__(self, exc_type, exc_value, traceback): return if numpyro.get_mask() is not False: + assert self.method is not None numpyro.factor( "_biased_corrected_log_likelihood", self.method(self.likelihoods, self.params, self.gibbs_state), @@ -740,10 +613,6 @@ def process_message(self, msg): if self.params is None: return - if msg["type"] == "_gibbs_state": - self.gibbs_state = msg["value"] - return - if msg["type"] == "sample" and msg["is_observed"]: assert msg["name"] not in self.params # store the likelihood for the estimator diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py index ecda7ddbf..a8b781ffb 100644 --- a/numpyro/infer/mcmc.py +++ b/numpyro/infer/mcmc.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod +from collections.abc import Callable from functools import partial from operator import attrgetter import os @@ -13,6 +14,7 @@ from jax import device_get, jit, lax, local_device_count, pmap, random, vmap import jax.numpy as jnp +from numpyro._typing import ModelArgs, ModelKwargs, StateT from numpyro.diagnostics import print_summary from numpyro.util import ( _is_under_jax_transform, @@ -158,6 +160,67 @@ def get_diagnostics_str(self, state): """ return "" + def refresh( + self, + state: StateT, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> StateT: + """ + Recompute every value cached in `state` that depends on `(model_args, model_kwargs)` + without advancing the chain, for example the potential energy and its gradient at the + current sample. Composite kernels call this before :meth:`sample` whenever the values + the target is conditioned on have changed. Kernels that cache nothing return `state`. + + The default raises `NotImplementedError`; only kernels that override it can be used as + blocks of :class:`~numpyro.infer.gibbs.Gibbs`. This is deliberate: a kernel that binds + its potential at `init` (e.g. :class:`~numpyro.infer.barker.BarkerMH`) would silently + target the wrong conditional if the default were the identity. + + :param state: current kernel state. + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model, including any + conditioning values. + :return: a state of the same type with refreshed cached values. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement `refresh` and cannot be used " + "as a block of a composite kernel." + ) + + def wrap_model(self, wrapper: Callable) -> "MCMCKernel": + """ + Return a copy of this kernel whose model is `wrapper(self.model)`. Kernels that hold + other kernels apply the wrapper recursively; kernels without a model return `self`. + Any function built lazily from the old model (potential, postprocess, sampler closures) + must be reset in the copy. The default raises `NotImplementedError`. + + :param wrapper: callable mapping a model to a model with the same call signature. + :return: a new kernel bound to the wrapped model. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement `wrap_model`." + ) + + def get_constrain_fn( + self, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> Callable: + """ + Return a function mapping the values in `state.` to the constrained + values a sibling kernel should be conditioned on. The default delegates to + :meth:`postprocess_fn`, which is always correct but may replay the model; kernels that + can constrain with transforms alone (for example :class:`~numpyro.infer.hmc.HMC`) + override it with a cheaper function. Composite kernels keep only the entries for the + sites the block owns. + + :param tuple model_args: arguments provided to the model. + :param dict model_kwargs: keyword arguments provided to the model. + :return: a callable mapping site values to constrained site values. + """ + return self.postprocess_fn(model_args, model_kwargs) + def _get_progbar_desc_str(num_warmup, phase, i): if phase is not None: @@ -721,7 +784,7 @@ def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs): states, last_state = _laxmap(partial_map_fn, map_args) elif self.chain_method == "parallel": states, last_state = pmap(partial_map_fn)(map_args) - elif callable(self.chain_method): + elif not isinstance(self.chain_method, str): states, last_state = self.chain_method(partial_map_fn)(map_args) else: assert self.chain_method == "vectorized" @@ -754,11 +817,9 @@ def get_samples(self, group_by_chain=False): samples = predictive(rng_key1, *model_args, **model_kwargs) """ - return ( - self._states[self._sample_field] - if group_by_chain - else self._get_states_flat()[self._sample_field] - ) + states = self._states if group_by_chain else self._get_states_flat() + assert states is not None, "`run` must be called before `get_samples`." + return states[self._sample_field] def get_extra_fields(self, group_by_chain=False): """ @@ -770,6 +831,7 @@ def get_extra_fields(self, group_by_chain=False): `extra_fields` keyword of :meth:`run`. """ states = self._states if group_by_chain else self._get_states_flat() + assert states is not None, "`run` must be called before `get_extra_fields`." return {k: v for k, v in states.items() if k != self._sample_field} def print_summary(self, prob=0.9, exclude_deterministic=True): @@ -781,7 +843,9 @@ def print_summary(self, prob=0.9, exclude_deterministic=True): at deterministic sites. """ # Exclude deterministic sites by default - sites = self._states[self._sample_field] + states = self._states + assert states is not None, "`run` must be called before `print_summary`." + sites = states[self._sample_field] if isinstance(sites, dict) and exclude_deterministic: state_sample_field = attrgetter(self._sample_field)(self._last_state) # Note: there might be the case that state.z is not a dictionary but @@ -792,7 +856,7 @@ def print_summary(self, prob=0.9, exclude_deterministic=True): if isinstance(state_sample_field, dict): sites = { k: v - for k, v in self._states[self._sample_field].items() + for k, v in states[self._sample_field].items() if k in state_sample_field } print_summary(sites, prob=prob) diff --git a/numpyro/infer/mixed_hmc.py b/numpyro/infer/mixed_hmc.py index a55070bc6..aaa5e29b7 100644 --- a/numpyro/infer/mixed_hmc.py +++ b/numpyro/infer/mixed_hmc.py @@ -2,23 +2,39 @@ # SPDX-License-Identifier: Apache-2.0 from collections import namedtuple +import copy from functools import partial +from typing import Any -from jax import grad, jacfwd, lax, random -from jax.flatten_util import ravel_pytree +import numpy as np + +import jax +from jax import random import jax.numpy as jnp -from numpyro.infer.hmc import momentum_generator -from numpyro.infer.hmc_gibbs import DiscreteHMCGibbs +from numpyro._typing import ConstrainFn, ModelArgs, ModelKwargs, ModelT, SiteValues +from numpyro.infer.gibbs import ( + ModelWrapper, + _flat_support_sizes, + conditioned, + discrete_latent_sites, + prototype_trace, + select_discrete_proposal, + with_conditioning, +) +from numpyro.infer.hmc import HMC, NUTS, momentum_generator from numpyro.infer.hmc_util import euclidean_kinetic_energy, warmup_adapter +from numpyro.infer.mcmc import MCMCKernel from numpyro.util import cond, fori_loop, identity MixedHMCState = namedtuple("MixedHMCState", "z, hmc_state, rng_key, accept_prob") -class MixedHMC(DiscreteHMCGibbs): +class MixedHMC(MCMCKernel): """ - Implementation of Mixed Hamiltonian Monte Carlo (reference [1]). + Implementation of Mixed Hamiltonian Monte Carlo (reference [1]). An HMC-family wrapper + that interleaves discrete updates inside the HMC trajectory; it can also be used as a + block of :class:`~numpyro.infer.gibbs.Gibbs`. .. note:: The number of discrete sites to update at each MCMC iteration (`n_D` in reference [1]) is fixed at value 1. @@ -68,31 +84,113 @@ class MixedHMC(DiscreteHMCGibbs): >>> assert abs(jnp.var(samples["x"]) - 4.36) < 0.5 """ + sample_field: str = "z" + def __init__( self, - inner_kernel, + inner_kernel: HMC, *, - num_discrete_updates=None, - random_walk=False, - modified=False, - ): - super().__init__(inner_kernel, random_walk=random_walk, modified=modified) - if inner_kernel._algo == "NUTS": + num_discrete_updates: int | None = None, + random_walk: bool = False, + modified: bool = False, + ) -> None: + if not isinstance(inner_kernel, HMC): + raise ValueError("inner_kernel must be an HMC sampler.") + if isinstance(inner_kernel, NUTS): raise ValueError( "The algorithm only works with HMC and and does not support NUTS." ) + if inner_kernel.model is None: + raise ValueError( + "MixedHMC does not support models specified via a potential function." + ) + self._model = inner_kernel.model + self.inner_kernel = inner_kernel.wrap_model(conditioned) self._num_discrete_updates = num_discrete_updates + self._random_walk = random_walk + self._modified = modified + self._discrete_proposal_fn = select_discrete_proposal(random_walk, modified) + # static metadata resolved at `init` + self._gibbs_sites: tuple[str, ...] = () + self._support_sizes_flat: np.ndarray | None = None + self._num_warmup = None + self._wa_update = None + + @property + def model(self) -> ModelT: + return self._model + + def get_diagnostics_str(self, state: MixedHMCState) -> str: + return self.inner_kernel.get_diagnostics_str(state.hmc_state) + + def _split(self, z: SiteValues) -> tuple[SiteValues, SiteValues]: + z_discrete = {k: v for k, v in z.items() if k in self._gibbs_sites} + z_hmc = {k: v for k, v in z.items() if k not in self._gibbs_sites} + return z_discrete, z_hmc + + def postprocess_fn( + self, model_args: ModelArgs, model_kwargs: ModelKwargs | None + ) -> ConstrainFn: + def fn(z: SiteValues) -> SiteValues: + z_discrete, z_hmc = self._split(z) + z_hmc = self.inner_kernel.postprocess_fn( + model_args, with_conditioning(model_kwargs, z_discrete) + )(z_hmc) + return {**z_discrete, **z_hmc} + + return fn - def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs): - rng_key, rng_r = random.split(rng_key) - state = super().init(rng_key, num_warmup, init_params, model_args, model_kwargs) - self._support_sizes_flat, _ = ravel_pytree( - {k: self._support_sizes[k] for k in self._gibbs_sites} + def get_constrain_fn( + self, model_args: ModelArgs, model_kwargs: ModelKwargs | None + ) -> ConstrainFn: + def fn(z: SiteValues) -> SiteValues: + z_discrete, z_hmc = self._split(z) + z_hmc = self.inner_kernel.get_constrain_fn( + model_args, with_conditioning(model_kwargs, z_discrete) + )(z_hmc) + return {**z_discrete, **z_hmc} + + return fn + + def init( + self, + rng_key: jax.Array, + num_warmup: int, + init_params: SiteValues | None, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> MixedHMCState: + model_kwargs = {} if model_kwargs is None else dict(model_kwargs) + rng_key, key_u, key_z, rng_r = random.split(rng_key, 4) + model_trace = prototype_trace( + conditioned(self._model), key_u, model_args, model_kwargs ) + self._gibbs_sites = discrete_latent_sites(model_trace) + if not self._gibbs_sites: + raise ValueError( + "Cannot detect any discrete latent variables in the model." + ) + self._support_sizes_flat = _flat_support_sizes(model_trace, self._gibbs_sites) if self._num_discrete_updates is None: self._num_discrete_updates = self._support_sizes_flat.shape[0] self._num_warmup = num_warmup + init_params = None if init_params is None else dict(init_params) + z_discrete = {} + for name in self._gibbs_sites: + if init_params and name in init_params: + z_discrete[name] = init_params.pop(name) + else: + z_discrete[name] = model_trace[name]["value"] + z_discrete = {k: jnp.asarray(v) for k, v in z_discrete.items()} + hmc_state = self.inner_kernel.init( + key_z, + num_warmup, + init_params or None, + model_args, + with_conditioning(model_kwargs, z_discrete), + ) + # NB: the warmup adaptation can not be performed in sub-trajectories (i.e. the hmc trajectory # between two discrete updates), so we will do it here, at the end of each MixedHMC step. _, self._wa_update = warmup_adapter( @@ -107,19 +205,46 @@ def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs): # In HMC, when `hmc_state.r` is not None, we will skip drawing a random momentum at the # beginning of an HMC step. The reason is we need to maintain `r` between each sub-trajectories. r = momentum_generator( - state.hmc_state.z, state.hmc_state.adapt_state.mass_matrix_sqrt, rng_r + hmc_state.z, hmc_state.adapt_state.mass_matrix_sqrt, rng_r ) - return MixedHMCState( - state.z, state.hmc_state._replace(r=r), state.rng_key, jnp.zeros(()) + z = {**z_discrete, **hmc_state.z} + return MixedHMCState(z, hmc_state._replace(r=r), rng_key, jnp.zeros(())) + + def refresh( + self, + state: MixedHMCState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> MixedHMCState: + """Delegates to the inner kernel with the current discrete values in the kwargs.""" + z_discrete, _ = self._split(state.z) + hmc_state = self.inner_kernel.refresh( + state.hmc_state, model_args, with_conditioning(model_kwargs, z_discrete) ) + return state._replace(hmc_state=hmc_state) - def sample(self, state, model_args, model_kwargs): + def wrap_model(self, wrapper: ModelWrapper) -> "MixedHMC": + kernel = copy.copy(self) + kernel._model = wrapper(self._model) + kernel.inner_kernel = self.inner_kernel.wrap_model(wrapper) + return kernel + + def sample( + self, + state: MixedHMCState, + model_args: ModelArgs, + model_kwargs: ModelKwargs | None, + ) -> MixedHMCState: model_kwargs = {} if model_kwargs is None else model_kwargs + assert self._support_sizes_flat is not None, ( + "`init` must be called before `sample`." + ) num_discretes = self._support_sizes_flat.shape[0] + support_sizes_flat = jnp.asarray(self._support_sizes_flat) def potential_fn(z_gibbs, z_hmc): - return self.inner_kernel._potential_fn_gen( - *model_args, _gibbs_sites=z_gibbs, **model_kwargs + return self.inner_kernel.get_potential_fn( + model_args, with_conditioning(model_kwargs, z_gibbs) )(z_hmc) def update_discrete( @@ -137,19 +262,31 @@ def update_discrete( hmc_state.potential_energy, partial(potential_fn, z_hmc=hmc_state.z), idx, - self._support_sizes_flat[idx], + support_sizes_flat[idx], ) # Algo 1, line 20: depending on reject or refract, we will update # the discrete variable and its corresponding kinetic energy. In case of # refract, we will need to update the potential energy and its grad w.r.t. hmc_state.z ke_discrete_i_new = ke_discrete[idx] + log_accept_ratio - grad_ = jacfwd if self.inner_kernel._forward_mode_differentiation else grad - z_discrete, pe, ke_discrete_i, z_grad = lax.cond( + + def refract(vals): + z_discrete_new, _, ke_discrete_i_new = vals + refreshed = self.inner_kernel.refresh( + hmc_state, + model_args, + with_conditioning(model_kwargs, z_discrete_new), + ) + return ( + z_discrete_new, + refreshed.potential_energy, + ke_discrete_i_new, + refreshed.z_grad, + ) + + z_discrete, pe, ke_discrete_i, z_grad = cond( ke_discrete_i_new > 0, (z_discrete_new, pe_new, ke_discrete_i_new), - lambda vals: ( - vals + (grad_(partial(potential_fn, vals[0]))(hmc_state.z),) - ), + refract, ( z_discrete, hmc_state.potential_energy, @@ -165,10 +302,8 @@ def update_discrete( return rng_key, hmc_state, z_discrete, ke_discrete, delta_pe_sum def update_continuous(hmc_state, z_discrete): - model_kwargs_ = model_kwargs.copy() - model_kwargs_["_gibbs_sites"] = z_discrete hmc_state_new = self.inner_kernel.sample( - hmc_state, model_args, model_kwargs_ + hmc_state, model_args, with_conditioning(model_kwargs, z_discrete) ) # each time a sub-trajectory is performed, we need to reset i and adapt_state @@ -217,7 +352,7 @@ def body_fn(i, vals): arrival_times, ) - z_discrete = {k: v for k, v in state.z.items() if k not in state.hmc_state.z} + z_discrete, _ = self._split(state.z) rng_key, rng_ke, rng_time, rng_r, rng_accept = random.split(state.rng_key, 5) # Algo 1, line 2: sample discrete kinetic energy ke_discrete = random.exponential(rng_ke, (num_discretes,)) @@ -226,9 +361,11 @@ def body_fn(i, vals): # the same job: the sub-trajectory length eta_t * M_t is the lag between two arrival time. arrival_times = random.uniform(rng_time, (num_discretes,)) # compute the amount of time to make `num_discrete_updates` discrete updates - total_time = (self._num_discrete_updates - 1) // num_discretes + jnp.sort( + num_discrete_updates = self._num_discrete_updates + assert num_discrete_updates is not None + total_time = (num_discrete_updates - 1) // num_discretes + jnp.sort( arrival_times - )[(self._num_discrete_updates - 1) % num_discretes] + )[(num_discrete_updates - 1) % num_discretes] # NB: total_time can be different from the HMC trajectory length, so we need to scale # the time unit so that total_time * time_unit = hmc_trajectory_length time_unit = state.hmc_state.trajectory_length / total_time @@ -273,7 +410,7 @@ def body_fn(i, vals): trajectory_length=hmc_state.trajectory_length ) hmc_state, z_discrete = cond( - random.bernoulli(rng_key, accept_prob), + random.bernoulli(rng_accept, accept_prob), (hmc_state_new, z_discrete_new), identity, (hmc_state, z_discrete), @@ -281,10 +418,12 @@ def body_fn(i, vals): ) # perform hmc adapting (similar to the implementation in hmc) + wa_update = self._wa_update + assert wa_update is not None adapt_state = cond( hmc_state.i < self._num_warmup, (hmc_state.i, accept_prob, (hmc_state.z,), hmc_state.adapt_state), - lambda args: self._wa_update(*args), + lambda args: wa_update(*args), hmc_state.adapt_state, identity, ) @@ -305,9 +444,7 @@ def body_fn(i, vals): z = {**z_discrete, **hmc_state.z} return MixedHMCState(z, hmc_state, rng_key, accept_prob) - def __getstate__(self): + def __getstate__(self) -> dict[str, Any]: state = self.__dict__.copy() state["_wa_update"] = None - state["_prototype_trace"] = None - state["_support_sizes_flat"] = None return state diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py index 41e28d945..64f35c7dc 100644 --- a/numpyro/infer/util.py +++ b/numpyro/infer/util.py @@ -5,7 +5,7 @@ from collections.abc import Sequence from contextlib import contextmanager from functools import partial -from typing import Callable, Optional +from typing import Callable, NamedTuple, Optional import warnings import numpy as np @@ -18,7 +18,7 @@ import numpyro from numpyro import distributions as dist -from numpyro._typing import TraceT +from numpyro._typing import ModelT, SiteValues, TraceT from numpyro.distributions import constraints from numpyro.distributions.transforms import biject_to from numpyro.distributions.util import is_identically_one, sum_rightmost @@ -330,6 +330,20 @@ def _unconstrain_reparam(params, site): return value +class _unconstrain_params(substitute): + """ + The handler :func:`potential_energy` uses to substitute unconstrained `params` through + :func:`_unconstrain_reparam`. Exposes `.params` so that model wrappers that need the + current unconstrained values (for example + :class:`~numpyro.infer.hmc_gibbs.estimate_likelihood`) can find it on the handler stack + with `isinstance` instead of inspecting `substitute_fn`. + """ + + def __init__(self, fn: ModelT, params: SiteValues) -> None: + super().__init__(fn, substitute_fn=partial(_unconstrain_reparam, params)) + self.params = params + + def potential_energy(model, model_args, model_kwargs, params, enum=False): """ (EXPERIMENTAL INTERFACE) Computes potential energy of a model given unconstrained params. @@ -348,9 +362,7 @@ def potential_energy(model, model_args, model_kwargs, params, enum=False): else: log_density_ = log_density - substituted_model = substitute( - model, substitute_fn=partial(_unconstrain_reparam, params) - ) + substituted_model = _unconstrain_params(model, params) # no param is needed for log_density computation because we already substitute log_joint, model_trace = log_density_( substituted_model, model_args, model_kwargs, {} @@ -508,12 +520,30 @@ def _find_valid_params(rng_key, exit_early=False): return (init_params, pe, z_grad), is_valid -def _get_model_transforms(model, model_args=(), model_kwargs=None): - model_kwargs = {} if model_kwargs is None else model_kwargs - model_trace = trace(model).get_trace(*model_args, **model_kwargs) +class _ModelTransforms(NamedTuple): + inv_transforms: dict + has_deterministic: bool + dynamic_support: bool + has_enumerate_support: bool + + +def _transforms_from_trace( + model_trace: TraceT, *, raise_warnings: bool = True +) -> _ModelTransforms: + """ + Inspect a model trace and collect the inverse transforms of its latent sample and param + sites together with the flags that decide how samples must be post-processed: + `has_deterministic` (the trace has `deterministic` sites, so constraining requires a + model replay to recover them), `dynamic_support` (a support depends on other values, so + constraining requires a model replay) and `has_enumerate_support` (the model has discrete + latent sites to enumerate). Does not mutate the trace. + + :param model_trace: a trace of the model. + :param bool raise_warnings: whether to emit the support and enumeration warnings. + """ inv_transforms = {} - # model code may need to be replayed in the presence of deterministic sites - replay_model = False + has_deterministic = False + dynamic_support = False has_enumerate_support = False for k, v in model_trace.items(): if v["type"] == "sample" and not v["is_observed"]: @@ -533,7 +563,7 @@ def _get_model_transforms(model, model_args=(), model_kwargs=None): f" enumerate support. But the {dist_name} distribution at" f" site {k} does not have enumerate support." ) - if enum_type is None: + if enum_type is None and raise_warnings: warnings.warn( "Some algorithms will automatically enumerate the discrete" f" latent site {k} of your model. In the future," @@ -544,7 +574,7 @@ def _get_model_transforms(model, model_args=(), model_kwargs=None): ) else: support = v["fn"].support - with helpful_support_errors(v, raise_warnings=True): + with helpful_support_errors(v, raise_warnings=raise_warnings): inv_transforms[k] = biject_to(support) # Note: the following code filters out most situations with dynamic supports args = () @@ -554,14 +584,65 @@ def _get_model_transforms(model, model_args=(), model_kwargs=None): args = ("lower_bound", "upper_bound") for arg in args: if not isinstance(getattr(support, arg), (int, float)): - replay_model = True + dynamic_support = True elif v["type"] == "param": - constraint = v["kwargs"].pop("constraint", constraints.real) - with helpful_support_errors(v, raise_warnings=True): + constraint = v["kwargs"].get("constraint", constraints.real) + with helpful_support_errors(v, raise_warnings=raise_warnings): inv_transforms[k] = biject_to(constraint) elif v["type"] == "deterministic": - replay_model = True - return inv_transforms, replay_model, has_enumerate_support, model_trace + has_deterministic = True + return _ModelTransforms( + inv_transforms, has_deterministic, dynamic_support, has_enumerate_support + ) + + +def _get_model_transforms(model, model_args=(), model_kwargs=None): + model_kwargs = {} if model_kwargs is None else model_kwargs + model_trace = trace(model).get_trace(*model_args, **model_kwargs) + info = _transforms_from_trace(model_trace) + # model code may need to be replayed in the presence of deterministic sites + replay_model = info.has_deterministic or info.dynamic_support + return info.inv_transforms, replay_model, info.has_enumerate_support, model_trace + + +def _prepare_model_for_potential( + model: ModelT, model_trace: TraceT, *, enum: bool +) -> ModelT: + """ + The model preparation :func:`initialize_model` performs before building a potential: + substitute `param`/`mutable` values from the trace, add a default PRNG key, wrap with + `enum(config_enumerate(...))` when `enum` is set, and validate plates. Shared with + :class:`~numpyro.infer.gibbs.DiscreteGibbs` so that it builds its potential from the same + prepared model as HMC. The wrapper order is relied upon by + :func:`find_valid_initial_params`. + + :param model: the model. + :param model_trace: a trace of `model`. + :param bool enum: whether to marginalize discrete latent sites by enumeration. + """ + # substitute param sites from model_trace to model so + # we don't need to generate again parameters of `numpyro.module` + model = substitute( + model, + data={ + k: site["value"] + for k, site in model_trace.items() + if site["type"] in ["param", "mutable"] + }, + ) + + model = _substitute_default_key(model) + + if enum: + from numpyro.contrib.funsor import config_enumerate, enum as enum_handler + + if not isinstance(model, enum_handler): + max_plate_nesting = _guess_max_plate_nesting(model_trace) + _validate_model(model_trace, plate_warning="error") + model = enum_handler(config_enumerate(model), -max_plate_nesting - 1) + else: + _validate_model(model_trace, plate_warning="loose") + return model def _partial_args_kwargs(fn, *args, **kwargs): @@ -729,18 +810,7 @@ def initialize_model( "`numpyro.deterministic` to add this value to the trace instead." ) - # substitute param sites from model_trace to model so - # we don't need to generate again parameters of `numpyro.module` - model = substitute( - model, - data={ - k: site["value"] - for k, site in model_trace.items() - if site["type"] in ["param", "mutable"] - }, - ) - - model = _substitute_default_key(model) + model = _prepare_model_for_potential(model, model_trace, enum=has_enumerate_support) constrained_values = { k: v["value"] @@ -750,16 +820,6 @@ def initialize_model( and not v["fn"].support.is_discrete } - if has_enumerate_support: - from numpyro.contrib.funsor import config_enumerate, enum - - if not isinstance(model, enum): - max_plate_nesting = _guess_max_plate_nesting(model_trace) - _validate_model(model_trace, plate_warning="error") - model = enum(config_enumerate(model), -max_plate_nesting - 1) - else: - _validate_model(model_trace, plate_warning="loose") - potential_fn, postprocess_fn = get_potential_fn( model, inv_transforms, diff --git a/numpyro/util.py b/numpyro/util.py index 0d1d54fac..ffda0b86d 100644 --- a/numpyro/util.py +++ b/numpyro/util.py @@ -142,7 +142,11 @@ def maybe_jit(fn: Callable, *args, **kwargs) -> Callable: def cond( - pred: bool, true_operand, true_fun: Callable, false_operand, false_fun: Callable + pred: bool | jax.Array, + true_operand, + true_fun: Callable, + false_operand, + false_fun: Callable, ) -> Any: if _DISABLE_CONTROL_FLOW_PRIM: if pred: @@ -849,9 +853,13 @@ def getter(obj): def _get_nested_attr(obj, field): """ - Helper function to recursively access attributes and dictionary keys. + Helper function to recursively access attributes, dictionary keys and, for tuples and + lists, decimal indices (e.g. ``"block_states.1.diverging"``). """ for attr in field.split("."): + if isinstance(obj, (tuple, list)) and attr.isdecimal(): + obj = obj[int(attr)] + continue try: obj = getattr(obj, attr) except AttributeError: diff --git a/pyproject.toml b/pyproject.toml index 06d504ffa..a1e9b0a61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -232,6 +232,7 @@ unresolved-import = "ignore" [tool.ty.src] include = [ + "numpyro/_typing.py", "numpyro/contrib/control_flow", "numpyro/contrib/funsor", "numpyro/contrib/hsgp", @@ -243,6 +244,11 @@ include = [ "numpyro/examples/datasets.py", "numpyro/handlers.py", "numpyro/infer/elbo.py", + "numpyro/infer/gibbs.py", + "numpyro/infer/hmc.py", + "numpyro/infer/hmc_gibbs.py", + "numpyro/infer/mcmc.py", + "numpyro/infer/mixed_hmc.py", "numpyro/infer/util.py", "numpyro/optim.py", "numpyro/primitives.py", diff --git a/test/infer/test_gibbs.py b/test/infer/test_gibbs.py new file mode 100644 index 000000000..4ad7fb85f --- /dev/null +++ b/test/infer/test_gibbs.py @@ -0,0 +1,735 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the composable Gibbs kernels in `numpyro.infer.gibbs`.""" + +from functools import partial +import pickle + +import numpy as np +from numpy.testing import assert_allclose +import pytest + +from jax import jit, pmap, random, vmap +import jax.numpy as jnp +from jax.scipy.linalg import cho_factor, cho_solve, solve_triangular + +import numpyro +from numpyro.contrib.control_flow import scan +import numpyro.distributions as dist +from numpyro.infer import ( + HMC, + MCMC, + NUTS, + BarkerMH, + CustomGibbs, + DiscreteGibbs, + DiscreteHMCGibbs, + Gibbs, + HMCGibbs, + MixedHMC, +) +from numpyro.infer.gibbs import ( + GIBBS_SITES_KWARG, + GibbsState, + any_changed, + conditioned, + discrete_latent_sites, + with_conditioning, +) +from numpyro.infer.hmc_gibbs import HMCGibbsState +from numpyro.util import cond, identity + + +def _linear_regression_gibbs_fn(X, XX, XY, Y, rng_key, gibbs_sites, hmc_sites): + N, P = X.shape + sigma = ( + jnp.exp(hmc_sites["log_sigma"]) + if "log_sigma" in hmc_sites + else hmc_sites["sigma"] + ) + sigma_sq = jnp.square(sigma) + covar_inv = XX / sigma_sq + jnp.eye(P) + L = cho_factor(covar_inv, lower=True)[0] + L_inv = solve_triangular(L, jnp.eye(P), lower=True) + loc = cho_solve((L, True), XY) / sigma_sq + beta_proposal = dist.MultivariateNormal(loc=loc, scale_tril=L_inv).sample(rng_key) + return {"beta": beta_proposal} + + +def _linear_data(seed, N, P, sigma): + np.random.seed(seed) + X = np.random.randn(N * P).reshape((N, P)) + XX = np.matmul(np.transpose(X), X) + Y = X[:, 0] + sigma * np.random.randn(N) + XY = np.sum(X * Y[:, None], axis=0) + return X, XX, XY, Y + + +def xy_model(): + x = numpyro.sample("x", dist.Normal(0.0, 2.0)) + y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + + +def xy_gibbs_fn(rng_key, gibbs_sites, hmc_sites): + y = hmc_sites["y"] + return {"x": dist.Normal(0.8 * (1 - y), jnp.sqrt(0.8)).sample(rng_key)} + + +def test_gibbs_helpers(): + model = conditioned(xy_model) + assert conditioned(model) is model + kwargs = with_conditioning({"a": 1, GIBBS_SITES_KWARG: {"x": 0.0}}, {"y": 1.0}) + assert kwargs == {"a": 1, GIBBS_SITES_KWARG: {"x": 0.0, "y": 1.0}} + with pytest.raises(ValueError, match="non-latent"): + with_conditioning({}, {"z": 1.0}, allowed=frozenset({"x"})) + assert not any_changed({"x": jnp.ones(2)}, {"x": jnp.ones(2)}) + assert any_changed({"x": jnp.ones(2), "c": 1}, {"x": jnp.ones(2), "c": 2}) + assert not any_changed({}, {}) + + +@pytest.mark.parametrize("kernel_cls", [HMC, NUTS]) +def test_linear_model_log_sigma( + kernel_cls, N=100, P=50, sigma=0.11, num_warmup=500, num_samples=500 +): + X, XX, XY, Y = _linear_data(0, N, P, sigma) + + def model(X, Y): + N, P = X.shape + log_sigma = numpyro.sample("log_sigma", dist.Normal(1.0)) + sigma = jnp.exp(log_sigma) + beta = numpyro.sample("beta", dist.Normal(jnp.zeros(P), jnp.ones(P))) + mean = jnp.sum(beta * X, axis=-1) + numpyro.deterministic("mean", mean) + numpyro.sample("obs", dist.Normal(mean, sigma), obs=Y) + + gibbs_fn = partial(_linear_regression_gibbs_fn, X, XX, XY, Y) + kernel = Gibbs([(CustomGibbs(gibbs_fn), ["beta"]), (kernel_cls(model), None)]) + mcmc = MCMC( + kernel, num_warmup=num_warmup, num_samples=num_samples, progress_bar=False + ) + mcmc.run(random.key(0), X, Y) + samples = mcmc.get_samples() + assert set(samples) == {"beta", "log_sigma", "mean"} + assert samples["mean"].shape == (num_samples, N) + beta_mean = np.mean(samples["beta"], axis=0) + assert_allclose(beta_mean, np.array([1.0] + [0.0] * (P - 1)), atol=0.05) + sigma_mean = np.exp(np.mean(samples["log_sigma"], axis=0)) + assert_allclose(sigma_mean, sigma, atol=0.25) + + +@pytest.mark.parametrize("kernel_cls", [HMC, NUTS]) +def test_linear_model_sigma( + kernel_cls, N=90, P=40, sigma=0.07, num_warmup=500, num_samples=500 +): + X, XX, XY, Y = _linear_data(1, N, P, sigma) + + def model(X, Y): + N, P = X.shape + sigma = numpyro.sample("sigma", dist.HalfCauchy(1.0)) + beta = numpyro.sample("beta", dist.Normal(jnp.zeros(P), jnp.ones(P))) + mean = jnp.sum(beta * X, axis=-1) + numpyro.sample("obs", dist.Normal(mean, sigma), obs=Y) + + gibbs_fn = partial(_linear_regression_gibbs_fn, X, XX, XY, Y) + kernel = Gibbs([(CustomGibbs(gibbs_fn), ["beta"]), (kernel_cls(model), None)]) + mcmc = MCMC( + kernel, num_warmup=num_warmup, num_samples=num_samples, progress_bar=False + ) + mcmc.run(random.key(0), X, Y) + beta_mean = np.mean(mcmc.get_samples()["beta"], axis=0) + assert_allclose(beta_mean, np.array([1.0] + [0.0] * (P - 1)), atol=0.05) + sigma_mean = np.mean(mcmc.get_samples()["sigma"], axis=0) + assert_allclose(sigma_mean, sigma, atol=0.25) + + +@pytest.mark.parametrize("kernel_cls", [HMC, NUTS]) +def test_gaussian_model(kernel_cls, D=2, num_warmup=5000, num_samples=5000): + np.random.seed(0) + cov = np.random.randn(4 * D * D).reshape((2 * D, 2 * D)) + cov = jnp.matmul(jnp.transpose(cov), cov) + 0.25 * jnp.eye(2 * D) + cov00 = cov[:D, :D] + cov01 = cov[:D, D:] + cov10 = cov[D:, :D] + cov11 = cov[D:, D:] + cov_01_cov11_inv = jnp.matmul(cov01, jnp.linalg.inv(cov11)) + cov_10_cov00_inv = jnp.matmul(cov10, jnp.linalg.inv(cov00)) + posterior_cov0 = cov00 - jnp.matmul(cov_01_cov11_inv, cov10) + posterior_cov1 = cov11 - jnp.matmul(cov_10_cov00_inv, cov01) + + def model(): + numpyro.sample( + "x", dist.MultivariateNormal(jnp.zeros(2 * D), covariance_matrix=cov) + ) + + def gaussian_gibbs_fn(rng_key, hmc_sites, gibbs_sites): + x1 = hmc_sites["x1"] + posterior_loc0 = jnp.matmul(cov_01_cov11_inv, x1) + x0_proposal = dist.MultivariateNormal( + loc=posterior_loc0, covariance_matrix=posterior_cov0 + ).sample(rng_key) + return {"x0": x0_proposal} + + def split_model(): + x0 = numpyro.sample( + "x0", dist.MultivariateNormal(jnp.zeros(D), covariance_matrix=cov00) + ) + numpyro.sample( + "x1", + dist.MultivariateNormal( + jnp.matmul(cov_10_cov00_inv, x0), covariance_matrix=posterior_cov1 + ), + ) + + kernel = Gibbs( + [ + (CustomGibbs(gaussian_gibbs_fn), ["x0"]), + (kernel_cls(split_model, dense_mass=True), None), + ] + ) + mcmc = MCMC( + kernel, num_warmup=num_warmup, num_samples=num_samples, progress_bar=False + ) + mcmc.run(random.key(0)) + x0_mean = np.mean(mcmc.get_samples()["x0"], axis=0) + x1_mean = np.mean(mcmc.get_samples()["x1"], axis=0) + x0_std = np.std(mcmc.get_samples()["x0"], axis=0) + x1_std = np.std(mcmc.get_samples()["x1"], axis=0) + assert_allclose(x0_mean, np.zeros(D), atol=0.25) + assert_allclose(x1_mean, np.zeros(D), atol=0.25) + assert_allclose(x0_std, np.sqrt(np.diagonal(cov00)), rtol=0.05) + assert_allclose(x1_std, np.sqrt(np.diagonal(cov11)), rtol=0.1) + + +def test_matches_hmc_gibbs_exactly(): + kernel = Gibbs([(CustomGibbs(xy_gibbs_fn), ["x"]), (NUTS(xy_model), None)]) + mcmc = MCMC(kernel, num_warmup=50, num_samples=50, progress_bar=False) + mcmc.run(random.key(0)) + ref = MCMC( + HMCGibbs(NUTS(xy_model), xy_gibbs_fn, ["x"]), + num_warmup=50, + num_samples=50, + progress_bar=False, + ) + ref.run(random.key(0)) + for name in ("x", "y"): + assert_allclose(mcmc.get_samples()[name], ref.get_samples()[name]) + assert isinstance(ref.last_state, HMCGibbsState) + assert ref.last_state.hmc_state is ref.last_state.block_states[-1] + assert ref.sampler.inner_kernel is ref.sampler.blocks[1][0] + assert ref.sampler.model is xy_model + + +def test_facade_extra_fields(): + def model(): + c = numpyro.sample("c", dist.Bernoulli(0.8)) + numpyro.sample("x", dist.Normal(c, 1.0)) + + mcmc = MCMC( + DiscreteHMCGibbs(NUTS(model)), num_warmup=20, num_samples=20, progress_bar=False + ) + mcmc.run( + random.key(0), + extra_fields=("hmc_state.potential_energy", "block_states.1.diverging"), + ) + extra = mcmc.get_extra_fields() + assert extra["hmc_state.potential_energy"].shape == (20,) + assert extra["block_states.1.diverging"].shape == (20,) + assert "acc. prob" in mcmc.sampler.get_diagnostics_str(mcmc.last_state) + + +def test_mixed_hmc_as_block(): + def model(): + c = numpyro.sample("c", dist.Bernoulli(0.8)) + x = numpyro.sample("x", dist.Normal(c, 1.0)) + y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + + def y_gibbs_fn(rng_key, gibbs_sites, hmc_sites): + x = hmc_sites["x"] + return {"y": dist.Normal(0.8 * (1 - x), jnp.sqrt(0.8)).sample(rng_key)} + + kernel = Gibbs( + [ + (CustomGibbs(y_gibbs_fn), ["y"]), + (MixedHMC(HMC(model, trajectory_length=1.2), num_discrete_updates=2), None), + ] + ) + mcmc = MCMC(kernel, num_warmup=500, num_samples=5000, progress_bar=False) + mcmc.run(random.key(0)) + samples = mcmc.get_samples() + assert set(samples) == {"c", "x", "y"} + ref = MCMC( + DiscreteHMCGibbs(NUTS(model)), + num_warmup=500, + num_samples=5000, + progress_bar=False, + ) + ref.run(random.key(0)) + for name in ("c", "x", "y"): + assert_allclose(samples[name].mean(), ref.get_samples()[name].mean(), atol=0.15) + + +def test_custom_gibbs_returned_keys_validated(): + def bad_gibbs_fn(rng_key, gibbs_sites, hmc_sites): + return {"z": jnp.zeros(())} + + kernel = Gibbs([(CustomGibbs(bad_gibbs_fn), ["x"]), (NUTS(xy_model), None)]) + mcmc = MCMC(kernel, num_warmup=2, num_samples=2, progress_bar=False) + with pytest.raises(ValueError, match="exactly the sites"): + mcmc.run(random.key(0)) + with pytest.raises(ValueError, match="initial values"): + CustomGibbs(bad_gibbs_fn).init(random.key(0), 1, None, (), {}) + with pytest.raises(ValueError, match="callable"): + CustomGibbs(None) + + +def test_nested_composite_with_deterministic_sites(): + def model(): + x = numpyro.sample("x", dist.Normal(0.0, 2.0)) + y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + z = numpyro.sample("z", dist.Normal(0.0, 2.0)) + numpyro.deterministic("s", x + y + z) + numpyro.sample("obs", dist.Normal(x + y + z, 1.0), obs=jnp.array([1.0])) + + def gibbs_fn(rng_key, gibbs_sites, hmc_sites): + # conditional of x given the other two sites: prior N(0, 4), likelihood N(1 - y - z, 1) + rest = hmc_sites["y"] + hmc_sites["z"] + assert set(hmc_sites) == {"y", "z"} + return {"x": dist.Normal(0.8 * (1 - rest), jnp.sqrt(0.8)).sample(rng_key)} + + inner = Gibbs([(CustomGibbs(gibbs_fn), ["x"]), (NUTS(model), ["y"])]) + kernel = Gibbs([(inner, ["x", "y"]), (NUTS(model), None)]) + mcmc = MCMC(kernel, num_warmup=500, num_samples=2000, progress_bar=False) + mcmc.run(random.key(0)) + samples = mcmc.get_samples() + assert set(samples) == {"x", "y", "z", "s"} + assert_allclose(samples["s"], samples["x"] + samples["y"] + samples["z"], rtol=1e-5) + # posterior of the sum: prior N(0, 12), obs 1 with unit variance + assert_allclose(samples["s"].mean(), 12 / 13, atol=0.15) + assert_allclose(samples["s"].std(), np.sqrt(12 / 13), rtol=0.15) + + +def test_three_hmc_blocks(): + def model(): + x = numpyro.sample("x", dist.Normal(0.0, 1.0)) + y = numpyro.sample("y", dist.HalfNormal(1.0)) + z = numpyro.sample("z", dist.Normal(0.0, 1.0)) + numpyro.sample("obs", dist.Normal(x + z, y), obs=jnp.array([0.5, -0.5, 1.0])) + + kernel = Gibbs([(NUTS(model), ["x"]), (NUTS(model), ["y"]), (NUTS(model), None)]) + mcmc = MCMC(kernel, num_warmup=500, num_samples=2000, progress_bar=False) + mcmc.run(random.key(0)) + ref = MCMC(NUTS(model), num_warmup=500, num_samples=2000, progress_bar=False) + ref.run(random.key(0)) + for name in ("x", "y", "z"): + assert_allclose( + mcmc.get_samples()[name].mean(), ref.get_samples()[name].mean(), atol=0.15 + ) + assert mcmc.get_samples()[name].std() > 0.2 + + +@pytest.mark.parametrize( + "blocks, match", + [ + ([], "at least one block"), + ( + [(NUTS(xy_model), ["x"]), (NUTS(xy_model), ["x", "y"])], + "more than one block", + ), + ([(NUTS(xy_model), ["x"])], "not owned by any block"), + ([(NUTS(xy_model), ["w"]), (NUTS(xy_model), None)], "not latent sample sites"), + ([(NUTS(xy_model), None), (NUTS(xy_model), None)], "At most one block"), + ([(BarkerMH(xy_model), None)], "does not implement `refresh`"), + ([(NUTS(xy_model), ["x"]), (NUTS(lambda: None), None)], "same model"), + ([(NUTS(potential_fn=lambda z: 0.0), None)], "potential function"), + ([(CustomGibbs(xy_gibbs_fn), ["x"])], "built on a model"), + ], +) +def test_invalid_blocks(blocks, match): + with pytest.raises(ValueError, match=match): + kernel = Gibbs(blocks) + kernel.init(random.key(0), 10, None, (), {}) + + +def test_init_errors(): + def discrete_model(): + c = numpyro.sample("c", dist.Bernoulli(0.3)) + numpyro.sample("x", dist.Normal(c, 1.0)) + + with pytest.raises(ValueError, match="Discrete latent sites"): + Gibbs([(NUTS(discrete_model), None)]).init(random.key(0), 10, None, (), {}) + + def subsample_model(data): + mean = numpyro.sample("mean", dist.Normal()) + with numpyro.plate("batch", data.shape[0], subsample_size=2): + numpyro.sample("obs", dist.Normal(mean, 1), obs=numpyro.subsample(data, 0)) + + with pytest.raises(ValueError, match="subsample plates"): + Gibbs([(NUTS(subsample_model), None)]).init( + random.key(0), 10, None, (jnp.ones(5),), {} + ) + + def dynamic_support_model(): + lb = numpyro.sample("lb", dist.Normal(0.0, 1.0)) + numpyro.sample("y", dist.Uniform(lb, lb + 1.0)) + + with pytest.raises(ValueError, match="value-dependent supports"): + Gibbs( + [(NUTS(dynamic_support_model), ["y"]), (NUTS(dynamic_support_model), None)] + ).init(random.key(0), 10, None, (), {}) + + kernel = Gibbs([(CustomGibbs(xy_gibbs_fn), ["x"]), (NUTS(xy_model), None)]) + with pytest.raises(ValueError, match="unknown sites"): + kernel.init(random.key(0), 10, {"w": jnp.zeros(())}, (), {}) + with pytest.raises(ValueError, match="single random key"): + kernel.init(random.split(random.key(0), 2), 10, None, (), {}) + + +def test_init_params_and_gated_refresh(): + kernel = Gibbs([(CustomGibbs(xy_gibbs_fn), ["x"]), (NUTS(xy_model), None)]) + init_params = {"x": jnp.array(0.25), "y": jnp.array(-0.75)} + state = kernel.init(random.key(0), 10, init_params, (), {}) + assert isinstance(state, GibbsState) + assert_allclose(state.z["x"], 0.25) + assert_allclose(state.z["y"], -0.75) + assert set(init_params) == {"x", "y"} + hmc_state = state.block_states[1] + # the HMC block's cached potential is consistent with the initial value of x + hmc = kernel.blocks[1][0] + refreshed = hmc.refresh(hmc_state, (), with_conditioning({}, {"x": state.z["x"]})) + assert_allclose(refreshed.potential_energy, hmc_state.potential_energy) + # a gated refresh under jit keeps the state structure + gated = jit( + lambda pred, s: cond( + pred, + s, + lambda s: hmc.refresh(s, (), {GIBBS_SITES_KWARG: {"x": 1.0}}), + s, + identity, + ) + ) + assert not jnp.allclose( + gated(True, hmc_state).potential_energy, + gated(False, hmc_state).potential_energy, + ) + + +def test_lifecycle_and_extra_fields(): + kernel = Gibbs([(CustomGibbs(xy_gibbs_fn), ["x"]), (NUTS(xy_model), None)]) + mcmc = MCMC(kernel, num_warmup=20, num_samples=20, progress_bar=False) + mcmc.warmup(random.key(0)) + mcmc.run( + random.key(1), + extra_fields=("block_states.1.diverging", "block_states.1.num_steps"), + ) + extra = mcmc.get_extra_fields() + assert extra["block_states.1.diverging"].shape == (20,) + assert extra["block_states.1.num_steps"].shape == (20,) + # a second run reuses the initialized kernel + mcmc.run(random.key(2)) + assert set(mcmc.get_samples()) == {"x", "y"} + # pickle then continue from the post warmup state + mcmc2 = pickle.loads(pickle.dumps(mcmc)) + mcmc2.post_warmup_state = mcmc2.last_state + mcmc2.run(random.key(3)) + assert set(mcmc2.get_samples()) == {"x", "y"} + assert mcmc2.sampler.get_diagnostics_str(mcmc2.last_state) + + +@pytest.mark.filterwarnings("ignore:There are not enough devices") +@pytest.mark.parametrize("chain_method", ["sequential", "parallel", vmap]) +def test_chain_methods(chain_method): + kernel = Gibbs([(CustomGibbs(xy_gibbs_fn), ["x"]), (NUTS(xy_model), None)]) + mcmc = MCMC( + kernel, + num_warmup=20, + num_samples=20, + num_chains=2, + chain_method=chain_method, + progress_bar=False, + ) + mcmc.run(random.key(0)) + mcmc.run(random.key(1)) + assert mcmc.get_samples(group_by_chain=True)["x"].shape == (2, 20) + + +def test_jit_model_args(): + def model(scale): + x = numpyro.sample("x", dist.Normal(0.0, scale)) + y = numpyro.sample("y", dist.Normal(0.0, scale)) + numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + + kernel = Gibbs([(CustomGibbs(xy_gibbs_fn), ["x"]), (NUTS(model), None)]) + mcmc = MCMC( + kernel, num_warmup=20, num_samples=20, progress_bar=False, jit_model_args=True + ) + mcmc.run(random.key(0), 2.0) + mcmc.run(random.key(1), 3.0) + assert set(mcmc.get_samples()) == {"x", "y"} + + +def test_scan_model(): + def model(T=5): + x0 = numpyro.sample("x0", dist.Normal(0.0, 1.0)) + sigma = numpyro.sample("sigma", dist.HalfNormal(1.0)) + + def transition(x, t): + x_new = numpyro.sample("x", dist.Normal(x, sigma)) + numpyro.sample("obs", dist.Normal(x_new, 0.5), obs=jnp.float32(t)) + return x_new, x_new + + scan(transition, x0, jnp.arange(T)) + + def sigma_gibbs_fn(rng_key, gibbs_sites, hmc_sites): + return {"sigma": dist.HalfNormal(1.0).sample(rng_key)} + + kernel = Gibbs([(CustomGibbs(sigma_gibbs_fn), ["sigma"]), (NUTS(model), None)]) + mcmc = MCMC(kernel, num_warmup=20, num_samples=20, progress_bar=False) + mcmc.run(random.key(0)) + assert mcmc.get_samples()["x"].shape == (20, 5) + + +def _discrete_blocks(model, inner_kernel=NUTS, **kwargs): + return Gibbs( + [ + (DiscreteGibbs(model, **kwargs), discrete_latent_sites), + (inner_kernel(model), None), + ] + ) + + +def _discrete_model(): + numpyro.sample("x", dist.Bernoulli(0.7).expand([3])) + numpyro.sample("y", dist.Binomial(10, 0.3)) + + +def test_discrete_gibbs_standalone(): + kernel = DiscreteGibbs(_discrete_model) + mcmc = MCMC(kernel, num_warmup=500, num_samples=5000, progress_bar=False) + mcmc.run(random.key(0)) + samples = mcmc.get_samples() + assert_allclose(jnp.mean(samples["x"], 0), 0.7 * jnp.ones(3), atol=0.05) + assert_allclose(jnp.mean(samples["y"], 0), 0.3 * 10, atol=0.1) + assert kernel._sites == ("x", "y") + assert_allclose(kernel._support_sizes_flat, np.array([2, 2, 2, 11])) + # refresh recomputes the potential energy at the current values + state = mcmc.last_state + pe = kernel.get_potential_fn((), {})(state.z) + assert_allclose(kernel.refresh(state, (), {}).potential_energy, pe, rtol=1e-5) + assert_allclose(state.potential_energy, pe, rtol=1e-5) + # a gated refresh under jit keeps the state structure + gated = jit( + lambda p, s: cond(p, s, lambda s: kernel.refresh(s, (), {}), s, identity) + ) + assert gated(True, state).potential_energy.dtype == state.potential_energy.dtype + # a second run re-initializes the kernel + mcmc.run(random.key(1)) + # pickle drops the prepared model only + kernel2 = pickle.loads(pickle.dumps(kernel)) + assert kernel2._prepared_model is None and kernel2._sites == ("x", "y") + + +def test_discrete_gibbs_errors(): + def mixed_model(): + c = numpyro.sample("c", dist.Bernoulli(0.8)) + numpyro.sample("x", dist.Normal(c, 1.0)) + + with pytest.raises(ValueError, match="cannot sample the latent sites"): + DiscreteGibbs(mixed_model).init(random.key(0), 10, None, (), {}) + with pytest.raises(ValueError, match="Cannot detect any discrete"): + DiscreteGibbs(xy_model).init(random.key(0), 10, None, (), {}) + with pytest.raises(RuntimeError, match="init"): + DiscreteGibbs(mixed_model).get_potential_fn() + + +def _mixed_discrete_model(): + c = numpyro.sample("c", dist.Bernoulli(0.7)) + z = numpyro.sample("z", dist.Normal(0.0, 1.0)) + numpyro.sample("obs", dist.Normal(z + c, 1.0), obs=jnp.array(1.0)) + + +@pytest.mark.parametrize( + "make_kernel, sites", + [ + (lambda: DiscreteGibbs(_discrete_model), {"x", "y"}), + ( + lambda: MixedHMC(HMC(_mixed_discrete_model, trajectory_length=1.2)), + {"c", "z"}, + ), + (lambda: DiscreteHMCGibbs(NUTS(_mixed_discrete_model)), {"c", "z"}), + ], + ids=["DiscreteGibbs", "MixedHMC", "DiscreteHMCGibbs"], +) +def test_discrete_kernels_init_under_pmap(make_kernel, sites): + # regression test for `init` under pmap's staging trace (multi-device CI runs + # chains with `chain_method="parallel"`): the flat support sizes stored on the + # kernel must stay static numpy, never a leaked tracer + kernel = make_kernel() + keys = random.split(random.key(0), 1) + states = pmap(lambda key: kernel.init(key, 10, None, (), {}))(keys) + discrete = ( + kernel if isinstance(kernel, (DiscreteGibbs, MixedHMC)) else kernel.blocks[0][0] + ) + assert isinstance(discrete._support_sizes_flat, np.ndarray) + states = pmap(lambda state: kernel.sample(state, (), {}))(states) + z = getattr(states, kernel.sample_field) + assert set(z) == sites + for value in z.values(): + assert value.shape[0] == 1 + assert np.isfinite(np.asarray(value, dtype=float)).all() + + +@pytest.mark.parametrize("num_chains", [1, 2]) +@pytest.mark.filterwarnings("ignore:There are not enough devices:UserWarning") +def test_discrete_gibbs_multiple_sites_chain(num_chains): + def model(): + numpyro.sample("x", dist.Bernoulli(0.7).expand([3])) + numpyro.sample("y", dist.Binomial(10, 0.3)) + + mcmc = MCMC( + _discrete_blocks(model), + num_warmup=1000, + num_samples=10000, + num_chains=num_chains, + progress_bar=False, + ) + mcmc.run(random.key(0)) + samples = mcmc.get_samples() + assert_allclose(jnp.mean(samples["x"], 0), 0.7 * jnp.ones(3), atol=0.05) + assert_allclose(jnp.mean(samples["y"], 0), 0.3 * 10, atol=0.1) + + +def test_discrete_gibbs_enum(): + def model(): + numpyro.sample("x", dist.Bernoulli(0.7), infer={"enumerate": "parallel"}) + y = numpyro.sample("y", dist.Binomial(10, 0.3)) + numpyro.deterministic("y2", y**2) + z = numpyro.sample("z", dist.Normal(0.0, 1.0)) + numpyro.sample("obs", dist.Normal(z + y, 1.0), obs=jnp.array(3.0)) + + kernel = _discrete_blocks(model) + mcmc = MCMC(kernel, num_warmup=1000, num_samples=10000, progress_bar=False) + mcmc.run(random.key(0)) + samples = mcmc.get_samples() + assert set(samples) == {"y", "y2", "z"} + assert kernel.blocks[0][0]._enum + # y | z has prior Binomial(10, 0.3) and likelihood N(3 - z, 1); the posterior mean of y + # is pulled towards 3 + assert 2.5 < jnp.mean(samples["y"]) < 3.5 + assert_allclose(samples["y2"], samples["y"] ** 2) + + +def test_discrete_gibbs_enum_potential_marginalizes(): + def model(): + x = numpyro.sample("x", dist.Bernoulli(0.7), infer={"enumerate": "parallel"}) + numpyro.sample("y", dist.Bernoulli(0.3)) + numpyro.sample("obs", dist.Normal(x, 1.0), obs=jnp.array(0.5)) + + kernel = DiscreteGibbs(model) + state = kernel.init(random.key(0), 10, None, (), {}) + pe = kernel.get_potential_fn((), {})({"y": jnp.array(1)}) + # exact marginal over x + likelihood = 0.7 * jnp.exp(dist.Normal(1.0, 1.0).log_prob(0.5)) + 0.3 * jnp.exp( + dist.Normal(0.0, 1.0).log_prob(0.5) + ) + expected = -(jnp.log(0.3) + jnp.log(likelihood)) + assert_allclose(pe, expected, rtol=1e-5) + assert state.z.keys() == {"y"} + + +@pytest.mark.parametrize("random_walk", [False, True]) +@pytest.mark.parametrize("modified", [False, True]) +def test_discrete_gibbs_bernoulli(random_walk, modified): + def model(): + numpyro.sample("c", dist.Bernoulli(0.8)) + + kernel = _discrete_blocks(model, random_walk=random_walk, modified=modified) + mcmc = MCMC(kernel, num_warmup=1000, num_samples=10000, progress_bar=False) + mcmc.run(random.key(0)) + samples = mcmc.get_samples()["c"] + assert_allclose(jnp.mean(samples), 0.8, atol=0.05) + + +def test_discrete_gibbs_improper_uniform(): + def model(): + numpyro.sample("c", dist.Bernoulli(0.8)) + numpyro.sample( + "u", dist.ImproperUniform(dist.constraints.unit_interval, (), ()) + ) + + mcmc = MCMC( + _discrete_blocks(model), num_warmup=10, num_samples=10, progress_bar=False + ) + mcmc.run(random.key(0)) + + +@pytest.mark.parametrize("modified", [False, True]) +def test_discrete_gibbs_gmm_1d(modified): + def model(probs, locs): + c = numpyro.sample("c", dist.Categorical(probs)) + numpyro.sample("x", dist.Normal(locs[c], 0.5)) + + probs = jnp.array([0.15, 0.3, 0.3, 0.25]) + locs = jnp.array([-2, 0, 2, 4]) + kernel = Gibbs( + [ + (DiscreteGibbs(model, modified=modified), discrete_latent_sites), + (NUTS(model, trajectory_length=1.2), None), + ] + ) + mcmc = MCMC(kernel, num_warmup=1000, num_samples=200000, progress_bar=False) + mcmc.run(random.key(0), probs, locs) + samples = mcmc.get_samples() + assert_allclose(jnp.mean(samples["x"]), 1.3, atol=0.1) + assert_allclose(jnp.var(samples["x"]), 4.36, atol=0.4) + assert_allclose(jnp.mean(samples["c"]), 1.65, atol=0.1) + assert_allclose(jnp.var(samples["c"]), 1.03, atol=0.1) + + +def test_three_blocks_doctest_model(): + def model(probs, locs): + c = numpyro.sample("c", dist.Categorical(probs)) + x = numpyro.sample("x", dist.Normal(locs[c], 0.5)) + y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + + def gibbs_fn(rng_key, gibbs_sites, hmc_sites): + x = hmc_sites["x"] + assert set(hmc_sites) == {"c", "x"} + return {"y": dist.Normal(0.8 * (1 - x), jnp.sqrt(0.8)).sample(rng_key)} + + kernel = Gibbs( + [ + (DiscreteGibbs(model), discrete_latent_sites), + (CustomGibbs(gibbs_fn), ["y"]), + (NUTS(model), None), + ] + ) + mcmc = MCMC(kernel, num_warmup=1000, num_samples=20000, progress_bar=False) + mcmc.run( + random.key(0), + jnp.array([0.15, 0.3, 0.3, 0.25]), + jnp.array([-2.0, 0.0, 2.0, 4.0]), + ) + samples = mcmc.get_samples() + assert set(samples) == {"c", "x", "y"} + + # reference: NUTS on the marginalized model + def ref_model(probs, locs): + x = numpyro.sample( + "x", dist.MixtureSameFamily(dist.Categorical(probs), dist.Normal(locs, 0.5)) + ) + y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + + ref = MCMC(NUTS(ref_model), num_warmup=1000, num_samples=20000, progress_bar=False) + ref.run( + random.key(0), + jnp.array([0.15, 0.3, 0.3, 0.25]), + jnp.array([-2.0, 0.0, 2.0, 4.0]), + ) + for name in ("x", "y"): + assert_allclose(samples[name].mean(), ref.get_samples()[name].mean(), atol=0.1) + assert_allclose(samples[name].std(), ref.get_samples()[name].std(), rtol=0.1) diff --git a/test/infer/test_hmc_gibbs.py b/test/infer/test_hmc_gibbs.py index f2077d2b0..544e8946b 100644 --- a/test/infer/test_hmc_gibbs.py +++ b/test/infer/test_hmc_gibbs.py @@ -13,7 +13,16 @@ import numpyro import numpyro.distributions as dist -from numpyro.infer import HMC, HMCECS, MCMC, NUTS, DiscreteHMCGibbs, HMCGibbs, MixedHMC +from numpyro.infer import ( + HMC, + HMCECS, + MCMC, + NUTS, + SA, + DiscreteHMCGibbs, + HMCGibbs, + MixedHMC, +) from numpyro.infer.util import log_density @@ -484,3 +493,82 @@ def gibbs_fn(rng_key, gibbs_sites, hmc_sites): mcmc.run(random.key(0)) samples = mcmc.get_samples() assert set(samples.keys()) == {"x", "y"} + + +def _subsample_model(data, subsample_size): + mean = numpyro.sample("mean", dist.Normal().expand((3,)).to_event(1)) + with numpyro.plate("batch", data.shape[0], dim=-1, subsample_size=subsample_size): + sub_data = numpyro.subsample(data, 1) + numpyro.sample("obs", dist.Normal(mean, 1).to_event(), obs=sub_data) + + +@pytest.mark.parametrize("use_proxy", [False, True]) +def test_hmcecs_lifecycle(use_proxy): + # regression test: the model used to be re-wrapped with `estimate_likelihood` at every + # `init`, so `warmup()` followed by `run()` raised on duplicated site names + true_loc = jnp.array([0.3, 0.1, 0.9]) + data = true_loc + dist.Normal(jnp.zeros(3), jnp.ones(3)).sample( + random.key(1), (1000,) + ) + proxy = HMCECS.taylor_proxy({"mean": true_loc}, degree=2) if use_proxy else None + kernel = HMCECS(NUTS(_subsample_model), proxy=proxy) + mcmc = MCMC(kernel, num_warmup=20, num_samples=20, progress_bar=False) + mcmc.warmup(random.key(0), data, 50) + mcmc.run(random.key(1), data, 50) + mcmc.run(random.key(2), data, 50) + assert mcmc.get_samples()["mean"].shape == (20, 3) + + +def test_hmc_gibbs_reuses_initialized_inner_kernel(): + def model(): + x = numpyro.sample("x", dist.Normal(0.0, 2.0)) + y = numpyro.sample("y", dist.Normal(0.0, 2.0)) + numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0])) + + def gibbs_fn(rng_key, gibbs_sites, hmc_sites): + y = hmc_sites["y"] + return {"x": dist.Normal(0.8 * (1 - y), jnp.sqrt(0.8)).sample(rng_key)} + + inner = NUTS(model) + MCMC(inner, num_warmup=10, num_samples=10, progress_bar=False).run(random.key(0)) + kernel = HMCGibbs(inner, gibbs_fn=gibbs_fn, gibbs_sites=["x"]) + init_params = {"x": jnp.array(0.5), "y": jnp.array(-0.5)} + mcmc = MCMC(kernel, num_warmup=10, num_samples=10, progress_bar=False) + mcmc.run(random.key(0), init_params=init_params) + assert set(mcmc.get_samples()) == {"x", "y"} + # `init_params` is not mutated by the kernel + assert set(init_params) == {"x", "y"} + # the wrapped copy did not leak into the user's kernel + assert inner.model is model + + +def test_hmcecs_conditioned_site_in_subsample_plate(): + def model(data): + mean = numpyro.sample("mean", dist.Normal()) + with numpyro.plate("batch", data.shape[0], subsample_size=2): + local = numpyro.sample("local", dist.Normal(mean, 1.0)) + numpyro.sample("obs", dist.Normal(local, 1), obs=numpyro.subsample(data, 0)) + + kernel = HMCECS(NUTS(model)) + with pytest.raises(ValueError, match="inside a subsample plate"): + kernel.init( + random.key(0), + 10, + None, + (jnp.ones(5),), + {"_gibbs_sites": {"local": jnp.ones(2)}}, + ) + + +def test_hmc_gibbs_public_names(): + from numpyro.infer.hmc_gibbs import ( # noqa: F401 + HMCECSState, + HMCGibbsState, + estimate_likelihood, + taylor_proxy, + ) + + with pytest.raises(ValueError, match="HMC or NUTS"): + HMCECS(SA(lambda: None)) + with pytest.raises(ValueError, match="potential function"): + HMCECS(NUTS(potential_fn=lambda z: 0.0)) diff --git a/test/infer/test_kernel_hooks.py b/test/infer/test_kernel_hooks.py new file mode 100644 index 000000000..be1826b02 --- /dev/null +++ b/test/infer/test_kernel_hooks.py @@ -0,0 +1,175 @@ +# Copyright Contributors to the Pyro project. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the composable-kernel hooks on `MCMCKernel` and the public `HMC` accessors.""" + +import pytest + +from jax import random, value_and_grad +import jax.numpy as jnp + +import numpyro +from numpyro import handlers +import numpyro.distributions as dist +from numpyro.infer import HMC, MCMC, NUTS +from numpyro.infer.mcmc import MCMCKernel +from numpyro.infer.util import ( + _prepare_model_for_potential, + _transforms_from_trace, + _unconstrain_params, + initialize_model, + potential_energy, +) +from numpyro.util import _get_nested_attr + + +def model(scale=1.0): + x = numpyro.sample("x", dist.Normal(0.0, scale)) + sigma = numpyro.sample("sigma", dist.HalfNormal(1.0)) + numpyro.deterministic("x2", x**2) + numpyro.sample("obs", dist.Normal(x, sigma), obs=jnp.array([0.3, -0.2])) + + +def dynamic_support_model(): + lb = numpyro.sample("lb", dist.Normal(0.0, 1.0)) + numpyro.sample("y", dist.Uniform(lb, lb + 1.0)) + + +def _init(kernel, *args, **kwargs): + return kernel.init(random.key(0), 10, None, model_args=args, model_kwargs=kwargs) + + +def test_mcmc_kernel_defaults_raise(): + class Dummy(MCMCKernel): + sample_field = "z" + + def init(self, rng_key, num_warmup, init_params, model_args, model_kwargs): + return None + + def sample(self, state, model_args, model_kwargs): + return state + + with pytest.raises(NotImplementedError, match="refresh"): + Dummy().refresh(None, (), {}) + with pytest.raises(NotImplementedError, match="wrap_model"): + Dummy().wrap_model(lambda m: m) + # the default constrain function is the (identity) postprocess function + assert Dummy().get_constrain_fn((), {})({"z": 1.0}) == {"z": 1.0} + + +@pytest.mark.parametrize("forward_mode", [False, True]) +def test_hmc_refresh_matches_value_and_grad(forward_mode): + kernel = NUTS(model, forward_mode_differentiation=forward_mode) + with pytest.raises(RuntimeError, match="init"): + kernel.get_potential_fn((2.0,), {}) + state = _init(kernel, 2.0) + # pretend the cached values are stale + stale = state._replace( + potential_energy=jnp.zeros(()), z_grad={k: 0 * v for k, v in state.z.items()} + ) + refreshed = kernel.refresh(stale, (2.0,), {}) + pe, z_grad = value_and_grad(kernel.get_potential_fn((2.0,), {}))(state.z) + assert jnp.allclose(refreshed.potential_energy, pe) + for k in z_grad: + assert jnp.allclose(refreshed.z_grad[k], z_grad[k]) + assert refreshed.energy is state.energy + # a different model argument gives a different potential + assert not jnp.allclose(kernel.refresh(stale, (0.1,), {}).potential_energy, pe) + + +def test_hmc_get_constrain_fn(): + kernel = HMC(model) + with pytest.raises(RuntimeError, match="init"): + kernel.get_constrain_fn((), {}) + state = _init(kernel) + transform_only = kernel.get_constrain_fn((), {})(state.z) + replay = kernel.get_constrain_fn((), {}, return_deterministic=True)(state.z) + assert set(transform_only) == {"x", "sigma"} + assert set(replay) == {"x", "sigma", "x2"} + for k in transform_only: + assert jnp.allclose(transform_only[k], replay[k]) + assert transform_only["sigma"] > 0 + assert jnp.allclose(replay["x2"], replay["x"] ** 2) + + +def test_hmc_get_constrain_fn_dynamic_support(): + kernel = HMC(dynamic_support_model) + state = _init(kernel) + assert kernel._dynamic_support + constrained = kernel.get_constrain_fn((), {})(state.z) + assert constrained["lb"] < constrained["y"] < constrained["lb"] + 1.0 + + +def test_hmc_wrap_model(): + kernel = NUTS(model) + _init(kernel, 2.0) + calls = [] + + def wrapper(m): + def wrapped(*args, **kwargs): + calls.append(1) + return m(*args, **kwargs) + + return wrapped + + wrapped = kernel.wrap_model(wrapper) + assert wrapped is not kernel + assert wrapped.model is not kernel.model + for attr in ("_init_fn", "_sample_fn", "_potential_fn_gen", "_postprocess_fn"): + assert getattr(wrapped, attr) is None + assert getattr(kernel, attr) is not None + _init(wrapped, 2.0) + assert calls + + with pytest.raises(ValueError, match="potential function"): + HMC(potential_fn=lambda z: z["x"] ** 2).wrap_model(wrapper) + + +def test_get_nested_attr_tuple_index(): + obj = {"a": (0, {"b": 3})} + assert _get_nested_attr(obj, "a.1.b") == 3 + assert _get_nested_attr(obj, "a.0") == 0 + + +def test_unconstrain_params_on_stack(): + found = {} + + class spy(handlers.Messenger): + def process_message(self, msg): + if msg["type"] == "sample" and msg["name"] == "obs": + for handler in numpyro.primitives._PYRO_STACK[::-1]: + if isinstance(handler, _unconstrain_params): + found["params"] = handler.params + + params = {"x": jnp.array(0.1), "sigma": jnp.array(-0.3)} + potential_energy(spy(model), (2.0,), {}, params) + assert found["params"] is params + + +def test_prepare_model_for_potential_matches_initialize_model(): + rng_key = random.key(1) + info = initialize_model(rng_key, model, model_args=(2.0,)) + prepared = _prepare_model_for_potential(model, info.model_trace, enum=False) + z = info.param_info.z + assert jnp.allclose(potential_energy(prepared, (2.0,), {}, z), info.potential_fn(z)) + + +def test_transforms_from_trace_flags(): + trace = handlers.trace(handlers.seed(model, random.key(0))).get_trace(2.0) + info = _transforms_from_trace(trace, raise_warnings=False) + assert set(info.inv_transforms) == {"x", "sigma"} + assert info.has_deterministic and not info.dynamic_support + assert not info.has_enumerate_support + trace = handlers.trace( + handlers.seed(dynamic_support_model, random.key(0)) + ).get_trace() + info = _transforms_from_trace(trace, raise_warnings=False) + assert info.dynamic_support and not info.has_deterministic + + +def test_mcmc_still_runs(): + mcmc = MCMC(NUTS(model), num_warmup=5, num_samples=5, progress_bar=False) + mcmc.run(random.key(0), 2.0) + assert set(mcmc.get_samples()) == {"x", "sigma", "x2"} + # re-init of an already-run kernel still works + mcmc.run(random.key(1), 2.0)