bayesmith.exact.gaussian
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
Reading a Gaussian off a node -- and checking the reading. The exact solves need three numbers a graph does not hand over directly: the prediction and its sigma (from an observed node's distribution), and the prior's centre and width (from a latent node's). ``dist_fn`` returns an opaque NumPyro distribution, so those numbers have to be *extracted* -- and, this package's rule being that a claim is checked rather than trusted, the extraction is checked too. **Two functions, deliberately.** :func:`gaussian_parts` is the fast path: one ``isinstance``, two attribute reads, fully traceable, and it is what runs inside ``jax.linearize`` on every solve. :func:`check_gaussian` is the guard: it probes the node's own ``log_prob`` at several points and refuses if the extracted ``(loc, scale)`` do not reproduce it. The guard uses Python floats and raises, so it cannot run inside a trace -- it runs once, on concrete values, when the block is built. Splitting them is not an optimisation, it is a correctness requirement: a guard that cannot run where the fast path runs must run *before* it, on the values the fast path will see. **Why the introspection is a fast path and not the answer.** Reading ``.loc``/``.scale`` off a ``Normal`` trusts the type. A ``Distribution`` subclass may override ``log_prob`` -- censored, tempered, or simply wrong -- and keep both attributes, so the type is evidence, not proof. The probe is what raises the bar, at a cost of five ``log_prob`` evaluations per node per block build. **What the probe does and does not establish.** It establishes that the extracted ``(loc, scale)`` reproduce the node's own density *at the probed points, elementwise, at the shape the node's value actually takes*. It does not establish agreement everywhere: a finite set of points cannot certify a claim about a function, and a correction shaped to vanish at exactly these offsets -- a quartic with roots there, say -- would pass while being wrong by hundreds of nats at the mode. That is an adversarial construction, and the threat this guard is placed against is accidental: a censored likelihood, a tempered one, a hand-written approximation. Those do not have roots at the probe points. Stated so the guarantee is not mistaken for a stronger one.
PROBE_OFFSETS unwrap gaussian_parts node_shape check_gaussian check_observed precision_parts observed_data_and_loc observation_parts noise_std_at precision_at
PROBE_OFFSETS
Defined in bayesmith.exact.gaussian, line 64.
PROBE_OFFSETS: tuple[float, ...] = (-3.0, -1.0, 0.0, 0.5, 2.0)unwrap
Defined in bayesmith.exact.gaussian, line 69.
unwrap(distribution: Any) -> AnyStrip wrappers that only change how ``log_prob`` is reduced. ``Independent`` is what ``.to_event(k)`` produces: the same per-element density with the last ``k`` batch dimensions summed. The elementwise reading underneath is what the exact solves want -- a diagonal covariance is diagonal whether or not its log-density was summed -- so the wrapper is removed rather than refused. ``ExpandedDistribution`` is what ``.expand(shape)`` produces, and it is stripped for the same reason: its ``log_prob`` is the base's ``log_prob`` broadcast to the larger batch shape, and its draws are independent copies, so its per-element density IS the base's. The base keeps its own smaller ``loc``, which is the "scalar loc, larger value" case a plated latent already has; :func:`node_shape` reads the value's shape off the unstripped distribution. Matched by exact type, as ``diagnose.structure`` does, so a subclass that changes the broadcast is not assumed away -- and the base is returned, never rebuilt, so :func:`check_gaussian` still probes the node's own density.
gaussian_parts
Defined in bayesmith.exact.gaussian, line 96.
gaussian_parts(graph: Graph, node: Node, env: dict[str, Any]) -> tuple[jax.Array, jax.Array]``(loc, scale)`` of ``node``'s distribution, ``scale`` broadcast to ``loc``.
Traceable and **unchecked** -- pair it with :func:`check_gaussian`, which
runs once on concrete values before any trace is opened.
A :class:`~bayesmith.distributions.ComplexNormal` is read the same way:
its ``loc`` is complex and its ``scale`` real, and the scale is the width
of EACH part. Everything downstream splits the latent into its real
degrees of freedom, so nothing here needs to know more than that -- see
:func:`~bayesmith.exact.block.real_parts`.
Raises:
NotGaussian: if the distribution is not a (possibly ``Independent``-
wrapped) ``Normal`` or ``ComplexNormal``. The *type* is static
under tracing, so this refusal is safe to make inside a trace. It
is a **classification outcome, not a fault**: P3b's dispatcher
catches it and routes the block to NUTS.
StructureError: if ``loc`` has an integer dtype. That one IS a fault --
a conjugate solve differentiates through ``loc``.node_shape
Defined in bayesmith.exact.gaussian, line 170.
node_shape(graph: Graph, node: Node, env: dict[str, Any]) -> tuple[int, ...]Shape of ``node``'s VALUE. Three sources, broadcast together: the distribution's own batch shape, the plate the node sits in, and -- for an observed node -- the data. All three are needed. A plated latent whose ``dist_fn`` takes no plated parent has a scalar ``loc`` and a plate-shaped value (the ordinary "N iid draws from one shared prior"); an unplated observed node conditioned on a vector has a scalar ``loc`` and a vector value. Taking any one source alone gets one of those two wrong. This must agree with what ``to_numpyro`` opens the site at, or the block's domain is a different space from the one NUTS samples -- ``test_node_shape_agrees_with_the_numpyro_bridge`` pins it. The distribution's batch and event shape are read off the node's own, UNSTRIPPED distribution as well as off ``loc``: ``.expand`` enlarges the batch shape while the base it wraps keeps its scalar ``loc``, so ``loc`` alone would report a scalar for an eight-element value.
check_gaussian
Defined in bayesmith.exact.gaussian, line 221.
check_gaussian(graph: Graph, node: Node, env: dict[str, Any], *, rtol: float | None=None) -> dict[float, float]Verify the extracted ``(loc, scale)`` really reproduce ``node``'s log_prob.
Costs ``len(PROBE_OFFSETS)`` evaluations of ``log_prob``. Runs on concrete
values, **outside** any trace.
**Probes at the shape the node's VALUE takes**, not at ``dist_fn``'s own
batch shape, and compares **elementwise**. Both matter, and both were
measured:
* A plated latent whose ``dist_fn`` takes no plated parent has a scalar
``loc`` and a plate-shaped value; so does an unplated observed node
conditioned on a vector. Probing at the scalar evaluates ``log_prob`` at
a shape the consumer never uses. Measured on this package's own
``plated_latent`` fixture pattern: a ``Distribution`` subclass correct
on a scalar and off by 1000 nats per element on an array passed the
guard with every reported error exactly ``0.0``, against a real
discrepancy of 2.0e6 nats. This is the same shape of defect P1 recorded
-- the guard and the thing it guards looking at different shapes -- and
it is fixed the same way.
* A summed comparison dilutes a localised defect by the magnitudes of the
correct entries. Measured: one wrong element out of 1e6, off by 50 nats,
reports a summed relative error of 1.95e-5 -- under the default rtol,
silently accepted -- and an elementwise error of 50.
Args:
graph, node, env: the node under test and the values its parents take.
rtol: tolerance on the relative disagreement, per element. Default
``1e3 * eps`` of ``loc``'s dtype, which leaves room for
accumulated roundoff without admitting a real difference in
density.
Returns:
``{offset: worst relative error over the node's entries}`` -- useful
for reporting how Gaussian a node is, not only whether it passes.
Raises:
NotGaussian: propagated from :func:`gaussian_parts`.
StructureError: if the scale is not strictly positive and finite, or
if any probe disagrees by more than ``rtol`` at any entry.check_observed
Defined in bayesmith.exact.gaussian, line 339.
check_observed(graph: Graph, node: Node, env: dict[str, Any], *, rtol: float | None=None) -> dict[Any, float]Probe an observed node's declared density, whichever covariance it has.
The eager guard the build path runs, and the router between the two:
=================== ==================================================
distribution what is run
=================== ==================================================
``Normal`` :func:`check_gaussian` -- unchanged, five offsets
against the ``(loc, scale)`` read off it
``CirculantNormal`` :func:`~bayesmith.exact.precision. check_positive_definite`, THEN
:func:`~bayesmith.exact.precision.check_precision`
=================== ==================================================
**The order is load-bearing, and it was measured.** ``check_precision``
does refuse an indefinite kernel on its own -- but only through NaN
propagation, and it names the wrong cause. Measured on three indefinite
kernels, including one whose entries are all positive: numpyro's own
``CirculantNormal.log_prob`` returns ``nan``, so ``check_precision``
reports ``linearity=nan, normalizer=nan``, and its message explains
``linearity`` as "the log-density is not quadratic, so it has no
covariance to extract". The log-density IS quadratic; the kernel simply
describes no realisable process. A user sent to look at their ``det``
nodes and ``linear_in`` declarations would be looking in the wrong place.
:func:`~bayesmith.exact.precision.check_positive_definite` says what is
actually wrong, so it goes first.
Refusing through a NaN is also fragile in a way that saying so is not: it
holds only for as long as ``log_prob`` keeps returning ``nan`` rather
than, say, clamping.
Args:
graph, node, env: the node under test and the values its parents take.
rtol: passed to whichever guard runs. Both default it the same way,
to ``1e3 * eps`` of ``loc``'s dtype.
Returns:
Whatever the guard that ran returns -- ``{offset: error}`` from
:func:`check_gaussian`, ``{"operator"/"normalizer"/"linearity":
error}`` from :func:`~bayesmith.exact.precision.check_precision`.
The two are different measurements and are deliberately not squashed
into one shape; every caller in this package discards them and reads
the exception instead.
Raises:
NotGaussian: for a distribution neither row accepts, exactly as
before -- still a classification outcome, not a fault.
StructureError: from :func:`check_gaussian`.
ValueError: from
:func:`~bayesmith.exact.precision.check_positive_definite`, when
the kernel is not a realisable autocovariance.
PrecisionMismatch: from
:func:`~bayesmith.exact.precision.check_precision`.precision_parts
Defined in bayesmith.exact.gaussian, line 418.
precision_parts(graph: Graph, node: Node, env: dict[str, Any]) -> tuple[jax.Array, Any]``(loc, Precision)`` -- the ONE object the whole exact path should read.
The widened counterpart of :func:`gaussian_parts`, which returns a
per-sample ``scale`` and so can only ever describe a diagonal covariance.
Traceable and **unchecked**; pair it with
:func:`~bayesmith.exact.precision.check_precision`, which verifies the
returned operator against the node's own density by the gradient identity
and runs eagerly, outside any trace.
The gate widens by exactly two rows and refuses everything else with the
same :class:`NotGaussian`, so a model this cannot express is still a
CLASSIFICATION outcome routed to NUTS rather than a fault:
================== ====================================================
distribution Precision
================== ====================================================
``Normal`` ``DiagonalPrecision`` -- unchanged behaviour
``CirculantNormal`` ``CirculantPrecision`` from ``.covariance_row``
================== ====================================================
``CirculantNormal`` is not a new idea imported from outside: numpyro ships
it, and its log-density agrees with ``CirculantPrecision``'s and with a
dense reference to 1.8e-15 -- three independent routes. The exact path was
the only thing refusing a correlated noise the rest of the stack could
already handle.
**``depends_on_prediction=True`` is accepted now.** ``45198f9`` refused
it, because ``fisher.py``'s variance-information term was derived for a
DIAGONAL ``N`` and nothing said it survived off-diagonal terms. It does,
in a form that covers both rows at once:
``1/2 tr(N^-1 d_a N N^-1 d_b N) = 1/2 sum_k d_a log lam_k d_b log lam_k``
whenever ``N``'s eigenbasis does not move with the parameters -- ``I`` for
a ``Normal``, the DFT for a ``CirculantNormal``. Derived in
``docs/derivations/variance_information_spectral.wls``, measured against a
dense finite-difference Fisher matrix in
``docs/probes/probe_9_correlated_variance_information.py``, and
implemented as
:func:`~bayesmith.exact.fisher._log_spectrum_curvature`, which
``fisher_information`` reaches through ``precision_of=``. The refusal was
written to be deleted and this is its deletion.
Raises:
NotGaussian: for any other distribution, as before.
StructureError: for an integer ``loc``, as before.observed_data_and_loc
Defined in bayesmith.exact.gaussian, line 522.
observed_data_and_loc(graph: Graph, env: dict[str, Any]) -> tuple[dict[str, jax.Array], dict[str, jax.Array]]``({obs: data}, {obs: loc})`` over every observed node, reading no covariance.
The two halves of :func:`observation_parts` that have nothing to do with
the noise. Split out because ``block.py`` wants exactly these -- ``isolate``
takes the ``loc`` and the block build takes the ``data`` -- and reading
them through ``observation_parts`` made both walk ``gaussian_parts``,
which refuses a correlated node. A block's design matrix and its data do
not depend on whether the samples are correlated.
Both are broadcast to :func:`node_shape`, so the dicts align leaf for leaf
and every reduction downstream is one ``jax.tree.map``.observation_parts
Defined in bayesmith.exact.gaussian, line 547.
observation_parts(graph: Graph, env: dict[str, Any]) -> tuple[dict[str, jax.Array], dict[str, jax.Array], dict[str, jax.Array]]``({obs: data}, {obs: loc}, {obs: scale})`` over every observed node.
All three are broadcast to :func:`node_shape`, so the three dicts align
leaf for leaf and every reduction downstream is one ``jax.tree.map``.
The DIAGONAL walk: ``scale`` comes from :func:`gaussian_parts`, so this
refuses a correlated node as it always did. The first two dicts are
:func:`observed_data_and_loc`, which does not, and callers wanting only
those should take them from there.noise_std_at
Defined in bayesmith.exact.gaussian, line 572.
noise_std_at(graph: Graph, values: dict[str, Any]) -> dict[str, jax.Array]``{obs: scale}`` with the latents at ``values`` -- the GLS seam.
``iterative_gls`` iterates this: solve at the current sigma, recompute
sigma at the new solution, repeat. Whether it moves at all is what
separates a single Wiener solve from a reweighting loop, and
``test_noise_std_at_moves_with_the_latent_only_for_a_prediction_dependent_node``
exercises both sides.precision_at
Defined in bayesmith.exact.gaussian, line 585.
precision_at(graph: Graph, values: dict[str, Any]) -> dict[str, Any]``{obs: Precision}`` with the latents at ``values`` -- the OPERATOR seam.
The counterpart of :func:`noise_std_at`, and deliberately **not** its
replacement. The seam SPLITS rather than renames, because the two sides
are different quantities and only one of them generalises:
============== ============================== =========================
seam what it produces who needs it
============== ============================== =========================
``precision_at`` ``N^-1`` as an operator the solve, the density,
the information
``noise_std_at`` per-sample sigma VALUES ``iterative_gls``'s
reweighting, the
linearity check's unit
============== ============================== =========================
``iterative_gls`` iterates sigma values -- solve at the current sigma,
recompute sigma at the new solution, repeat -- and an operator has none to
iterate. Renaming ``noise_std_at`` to this would have left that loop with
nothing to say, which is why B9's step 4 is a split and not the rename its
own text describes.
A ``Precision`` is what a correlated node can supply and a sigma dict is
not, so this is the producer a ``CirculantNormal`` observation reaches the
solver through. For a diagonal model the two agree by construction, and
``test_precision_at_agrees_with_diagonal_from_noise_std_at`` pins that:
this is not a second implementation of the diagonal case, it is
:func:`precision_parts` over the same nodes ``observation_parts`` walks.
Traceable and **unchecked**, like the parts it is built from. Pair it with
:func:`check_gaussian` (diagonal) or
:func:`~bayesmith.exact.precision.check_precision` (either), which run
eagerly at the boundary.