bayesmith.bridge.numpyro_bridge
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
Turning a graph into a NumPyro model, and running NUTS on it. This is the last row of the dispatch table: whatever structure bayesmith cannot solve exactly is handed to NumPyro. It is also the oracle every exact path is checked against, because a graph that qualifies for an exact method always also qualifies for NUTS.
to_numpyro
Defined in bayesmith.bridge.numpyro_bridge, line 40.
to_numpyro(graph: Graph | ReducedGraph) -> Callable[..., dict[str, Any]]Build a NumPyro model that declares the same joint distribution.
Latent and observed nodes become ``numpyro.sample`` sites carrying the
graph's own node names, so posterior samples come back keyed by them.
Deterministic nodes are recorded with ``numpyro.deterministic`` so they
appear in traces and predictives without contributing a density.
The model takes ONE optional argument, ``observed``, which is the
migration's stated bridge convention:
================== ==================================================
``observed`` what the observed nodes are conditioned on
================== ==================================================
``None`` (default) each node's own declared data -- today's behaviour,
so every existing caller reads unchanged
a mapping that node's entry, or ``None`` for a node the
mapping omits, which UNconditions it
``{}`` nothing at all: the prior predictive
================== ==================================================
**Why the unconditioned mode is worth having, measured.** ``Predictive``
over a model whose ``obs=`` is baked in returns the observed node's DATA,
identical in every draw -- standard deviation ``0`` across 3000 draws
against a prediction spread of 0.26. That is correct of NumPyro and is
not what "posterior predictive" usually means, so the mode that gives the
other reading has to exist rather than be assumed.init_to_declared
Defined in bayesmith.bridge.numpyro_bridge, line 201.
init_to_declared(graph: Graph | ReducedGraph) -> AnyA NumPyro init strategy that starts where the graph's priors sit.
Pass it as ``nuts(graph, key, nuts_options={"init_strategy":
init_to_declared(graph)})``.
**Not a tuning knob.** NumPyro's default is ``init_to_uniform``, which
draws in the UNCONSTRAINED space with no knowledge of where the graph's
priors are; a posterior far narrower than its prior is a needle,
``init_to_uniform`` lands in the haystack, and warmup then adapts a step
size for wherever it landed. Measured on a power law whose amplitude has a
prior of 1e4 and a posterior width of 0.4, two chains of 400:
============================================== ====== ======
init r_hat ESS
============================================== ====== ======
default (``init_to_uniform``) 1609 1.0
the declared values 1.006 138.6
============================================== ====== ======
``nuts``'s own docstring already recorded that measurement and named the
remedy; what was missing was the remedy itself, spelled once. Written by
hand, "the declared values" is a second statement of the prior centres,
and the whole reason
:func:`~bayesmith.dispatch.classify.prior_environment` is public is that a
second spelling of that point lets the sampler start somewhere the
classifier never looked.
The declared point does not have to be GOOD -- only somewhere a gradient
can be followed.nuts
Defined in bayesmith.bridge.numpyro_bridge, line 241.
nuts(graph: Graph | ReducedGraph, key: jax.Array, *, 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, _control: Any=None) -> dict[str, jax.Array]Sample the posterior of ``graph`` with NUTS.
``chain_method`` and ``nuts_options`` are here because
:meth:`~bayesmith.dispatch.plan.InferencePlan.sample` promises them on
every path, and two of its five shapes -- the graph with no exact block,
and the SNIS collapse -- run through this function. Until they existed
those two shapes silently ignored both keywords while the mixed shape,
which reaches ``HMCGibbs`` through
:func:`~bayesmith.exact.gibbs.assemble`, honoured them. The names and the
defaults are ``assemble``'s, so the two spellings of "run a chain" take
the same words.
Args:
graph: the model.
key: a PRNG key.
num_warmup: adaptation draws, discarded.
num_samples: retained draws per chain.
num_chains: independent chains.
chain_method: how ``num_chains`` are run -- ``"sequential"``,
``"parallel"`` or ``"vectorized"``. All three are numpyro's own
and all three are legal here; ``assemble`` refuses
``"vectorized"`` for a Gibbs sweep, but that refusal is about
``HMCGibbs.init`` and does not apply to a bare kernel.
progress_bar: whether NumPyro prints progress.
nuts_options: keywords for the ``NUTS`` kernel itself
(``target_accept_prob``, ``max_tree_depth``, ``dense_mass``, ...).
**``init_strategy`` goes here, and on a narrow posterior it is
not a tuning knob.** NumPyro's default is ``init_to_uniform``,
which draws in the UNCONSTRAINED space with no knowledge of
where the graph's priors sit; a posterior far narrower than its
prior is a needle, and warmup then adapts a step size for
wherever it landed in the haystack. Measured on a power law
whose amplitude has a prior of 1e4 and a posterior width of
0.4, two chains of 400:
============================================== ====== ======
init r_hat ESS
============================================== ====== ======
default (``init_to_uniform``) 1609 1.0
``init_to_value`` at the declared values 1.006 138.6
============================================== ====== ======
The declared point does not have to be good -- only somewhere a
gradient can be followed. This is rheplicant's
``init_to_declared`` lesson, and the lesson is what was carried:
the remedy already exists here as a keyword, so what was missing
was the sentence saying so.
Returns:
A mapping from latent node name to its draws.predict
Defined in bayesmith.bridge.numpyro_bridge, line 320.
predict(graph: Graph | ReducedGraph, samples: Mapping[str, Any], key: jax.Array | None=None) -> dict[str, jax.Array]Posterior predictive: every node's value over a stack of draws.
Args:
graph: the model the draws came from.
samples: ``{latent name: (n_draws, *latent shape)}`` -- what
:func:`nuts` returns, unchanged.
key: PRNG key for any node the draws do not fix. Fixed by default,
so a predictive is reproducible.
Returns:
``{node name: (n_draws, *node shape)}``, deterministic nodes
included. **An OBSERVED node comes back as its data**, identical in
every draw, because the model is conditioned on it -- measured, a
standard deviation of ``0`` across 3000 draws against a prediction
spread of 0.26. For the noiseless prediction over the same draws use
:func:`~bayesmith.exact.fisher.push_forward`; for a genuine
predictive, call this graph's model with ``observed={}`` so the
observed sites are drawn rather than conditioned.
Raises:
GraphError: if a latent is missing from ``samples``; if a stack's
PER-SAMPLE shape is not the latent's own; or if the stacks
disagree about how many draws there are.
Note:
**Why the per-sample shape is checked and not just the name.** This
is rheplicant's ``predict_from_samples`` guard, ported because the
failure it prevents is reachable here and is silent. Measured on a
length-3 latent with three draws: handing the stack in TRANSPOSED
returns a finite, correctly-shaped ``(3, 3)`` predictive whose every
entry is wrong --
correct [[0, 2, 6], [3, 8, 15], [6, 14, 24]]
transposed [[0, 6, 18], [1, 8, 21], [2, 10, 24]]
-- because NumPyro's ``Predictive`` maps over the leading axis and
has no independent statement of what the latent's shape is. The
graph does have one: each latent's ``dist_fn``. A non-square
transposition raises a broadcast error from three layers down that
names neither the site nor the axis; a square one raises nothing at
all.