bayesmith.exact.gibbs
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
The ``gibbs_fn`` numpyro's ``HMCGibbs`` calls, and the assembly around it. ``HMCGibbs``'s own docstring says "it is the user's responsibility to provide a correct implementation of ``gibbs_fn`` that samples from the corresponding posterior conditional". **That sentence is this package's product.** NumPyro supplies the scaffolding and cannot know which conditionals are exactly solvable, nor solve them; the graph is what says so, and :mod:`bayesmith.exact.block` and :mod:`bayesmith.exact.solve` are what do it. **Two sweeps, one factory.** ``method="gcr"`` draws the conditional exactly: sigma does not move with the block, so the frozen-sigma Gaussian *is* the conditional. ``method="gcr+mh"`` is for the case sigma does move -- then the frozen-sigma draw is only a **proposal**, and the sweep is an independence-proposal Metropolis step whose correctness argument is :func:`_mh_step`'s docstring and is the single largest correction in P3b's spec. **Why this module takes a graph and a name list rather than a plan object.** ``InferencePlan`` lives in :mod:`bayesmith.dispatch`, which is the layer *above* this one: ``dispatch`` reads ``exact``'s checkers, and **no module in ``exact`` imports ``dispatch`` at module scope**. Importing the plan here would invert that. So :func:`gibbs_factory` and :func:`assemble` take the primitives they actually need and the plan layer adapts to them -- a layering decision, not an oversight. Shared prior anchoring lives in :mod:`bayesmith.exact._environment`, so ``exact`` has no dependency on ``dispatch``, including function-scope imports. The full import boundary is checked in ``tests/test_layering.py``.
GIBBS_METHODS
Defined in bayesmith.exact.gibbs, line 63.
GIBBS_METHODS: tuple[str, ...] = ('gcr', 'gcr+mh')gibbs_factory
Defined in bayesmith.exact.gibbs, line 239.
gibbs_factory(graph: Graph, names: Iterable[str], *, tol: float, method: str='gcr', sigma_rebuild: bool=False, maxiter: int | None=None, precision: Mapping[str, Any] | None=None) -> Callable[..., dict[str, jax.Array]]Build the callable ``HMCGibbs`` invokes once per sweep.
Args:
graph: the model. ``log p`` inside the sweep is its own ``log_joint``,
so the Gibbs step and NUTS cannot target different posteriors.
names: the block's members -- the ``gibbs_sites``.
tol: CG tolerance, chosen at compile time as
``CONVERGENCE_TARGET / condition_bound(...)``. The in-sweep
convergence guard is **off**: it costs ``POWER_ITERATIONS`` extra
operator applications per solve and ``eqx.error_if`` cannot help a
sweep it cannot interrupt. So this number is the only thing
standing between the sweep and a silently over-narrow posterior,
since an unconverged CG leaves the prior-dominated directions at
their starting value. Never leave it at a default with the guard
off.
method: one of :data:`GIBBS_METHODS`. ``"gcr"`` for a genuine
conditional draw; ``"gcr+mh"`` when sigma depends on the block
itself and the frozen-sigma draw is only a PROPOSAL.
sigma_rebuild: recompute the covariance inside the sweep, at the
current ``at``. Required whenever any observed node's scale has a
latent ancestor -- including one OUTSIDE the block, which
:func:`~bayesmith.exact.gls.check_prediction_dependence` cannot
see because it only ever moves block members. Measured on
``contrast_sigma_pair`` split as block ``{a}`` with ``b`` outside
(sigma ``base * exp(a - b)``): the movement probe reads
``1.71828183e+00`` -- bitwise identical at ``b = 0.5`` and
``b = 2.0`` -- while hoisting sigma gives a posterior 1.64x and
7.35x too WIDE at those two. The criterion is therefore
structural, not numeric; a movement probe is precisely what
cannot see it.
maxiter: CG iteration cap, passed to every solve in the sweep.
precision: an x-independent ``{observed: N^-1}`` of the caller's own
choosing, used verbatim in place of everything above.
**Correctness does not depend on this being any good** -- see
:func:`_mh_step`; a worse covariance costs acceptance rate, not
validity. Useful when the GLS fixed point is too expensive to
recompute every sweep, or when a deliberately loose covariance is
wanted. Build one with
:func:`~bayesmith.exact.precision.diagonal_from` from a sigma
dict, or take :attr:`~bayesmith.exact.gls.GLSResult.precision`.
Returns:
A function taking ``rng_key``, ``gibbs_sites`` and ``hmc_sites``
**by keyword** -- that is numpyro's actual calling convention, see
``HMCGibbs.sample`` -- and returning exactly ``names``, in constrained
space. Returning a subset raises an ``AssertionError`` upstream whose
message is empty; returning an extra key naming a NUTS latent is
silently ignored. Both measured on numpyro 0.21.0.
Raises:
ValueError: if ``method`` is not in :data:`GIBBS_METHODS`.
NotGaussian: if the covariance must be hoisted
(``sigma_rebuild=False``, no ``precision``) and some latent has no
prior centre to hoist at.assemble
Defined in bayesmith.exact.gibbs, line 435.
assemble(graph: Graph, names: Iterable[str], *, tol: float, method: str='gcr', sigma_rebuild: bool=False, maxiter: int | None=None, precision: Mapping[str, Any] | None=None, num_warmup: int=1000, num_samples: int=2000, num_chains: int=1, chain_method: str='sequential', progress_bar: bool=False, nuts_options: Mapping[str, Any] | None=None) -> MCMCWire the block's sweep into ``HMCGibbs`` and hand back an unrun ``MCMC``.
Unrun on purpose: ``mcmc.run(key)`` is the caller's, so the plan layer
keeps control of the key and of whatever it wants to do with
``mcmc.get_extra_fields()`` afterwards.
Args:
graph, names, tol, method, sigma_rebuild, maxiter, precision: passed
straight to :func:`gibbs_factory`.
num_warmup, num_samples, num_chains, progress_bar: ``MCMC``'s own.
chain_method: ``"sequential"`` or ``"parallel"``. Both were run
end-to-end on this package's own fixtures.
nuts_options: keywords for the inner ``NUTS`` kernel
(``target_accept_prob``, ``max_tree_depth``, ``dense_mass``, ...).
Raises:
NotImplementedError: for ``chain_method="vectorized"``. Measured on
numpyro 0.21.0: ``HMC.init`` branches on ``rng_key.ndim`` but
``HMCGibbs.init`` does not, so it calls ``random.split`` on the
batched key MCMC hands it and raises ``ValueError: split accepts a
single key, but was given a key array of shape (2,) != ()`` --
before ``gibbs_fn`` is ever called. Refusing it here turns that
into a sentence, and pins the limitation so an upstream fix
surfaces as a decision rather than a behaviour change.
ValueError: if ``names`` covers every latent, leaving the inner NUTS
with nothing to sample.