bayesmith.exact.block
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
The affine action of a group of latents on every prediction.
No matrix is ever formed: ``A`` comes from ``jax.linearize`` and ``A^T`` from
``jax.vjp``, so a block with 10^6 degrees of freedom costs one forward
evaluation per application. That is what makes the conjugate-Gaussian solves
in :mod:`bayesmith.exact.solve` tractable at all.
**Where everything comes from.** rheplicant's ``unchecked_operator`` takes a
``ParameterSpace`` for the latents and their priors and a ``Pipeline`` for the
prediction. Here both come from the graph: the prediction is every observed
node's ``loc``, the prior is each member's own distribution, and sigma is
every observed node's ``scale``. There is consequently **no keyword to
override any of them** -- and so none of rheplicant's reconciliation
machinery (``_reconcile``, ``_agrees``, ``_resolve_prior``,
``_require_prior_std``) has anything to reconcile. One statement of the prior
means the exact solve and NUTS cannot target different posteriors.
**One spelling.** The domain is ``{latent: array}`` and the codomain is
``{observed: array}``, always, for one member or several. rheplicant carries
a ``name=``/``names=`` dual spelling plus ``as_dict`` to paper over it; the
graph makes that unnecessary because everything already has a name.LinearBlock is_complex real_parts domain_zero domain_centre variance_parts largest_variance isolate unchecked_operator
LinearBlock
Defined in bayesmith.exact.block, line 64.
class LinearBlock
names: tuple[str, ...]
shape: dict[str, tuple[int, ...]]
dtype: dict[str, Any]
offset: dict[str, jax.Array]
forward: Callable[[Any], dict[str, jax.Array]]
adjoint: Callable[[Any], dict[str, jax.Array]]
data: dict[str, jax.Array]
prior_mean: dict[str, jax.Array]
prior_std: dict[str, jax.Array]``A x + offset``: what a group of latents does to every prediction.
Deliberately a plain frozen dataclass rather than an ``eqx.Module``: this
is a derived linear-algebra handle, not a differentiable model.
``forward`` and ``adjoint`` close over a traced computation, so a block is
something you build where you need it, not a pytree to carry around.
Attributes:
names: the block's members, in the caller's order.
shape, dtype: ``{name: ...}`` describing the domain.
offset: ``{observed: prediction at x = 0}`` -- everything OUTSIDE the
block contributes.
forward: ``{name: x} -> {observed: A x}``, from ``jax.linearize``.
adjoint: ``{observed: y} -> {name: A^T y}``, from ``jax.vjp``.
data: ``{observed: value}``, broadcast to the prediction's shape.
prior_mean, prior_std: ``{name: ...}``, read off each member's own
distribution at ``at``.
Adjoint convention: ``adjoint`` is exactly ``jax.vjp``, so the identity
that holds is over the **real** inner product::
sum(x * adjoint(y)) == sum(forward(x) * y)
which is the pairing a Gaussian likelihood forms.
``test_adjoint_is_the_transpose_under_the_real_inner_product`` pins both
halves so the distinction cannot rot into a silent factor.is_complex
Defined in bayesmith.exact.block, line 104.
is_complex(dtype: Any) -> boolWhether a declared latent dtype is complex.
real_parts
Defined in bayesmith.exact.block, line 109.
real_parts(block: LinearBlock) -> tuple[Callable, Callable]Convert between the block's DOMAIN and its real degrees of freedom.
A complex member is carried as ``(real, imag)``; a real one is left alone,
NOT wrapped in a one-element tuple -- there would be nothing to unwrap it
back from, and a uniform wrapper would only move the asymmetry somewhere
less visible.
Not bookkeeping pedantry. Every prediction here is real, so the map from
complex coefficients to data is **R-linear but not C-linear**, and a
Krylov method run over C would be minimising a different objective. The
split makes the vector space the one the posterior actually lives on.
It also keeps this module's adjoint story true rather than merely
convenient. :func:`~bayesmith.exact.solve.normal_operator` and
``_conjugate_solve``'s ``pair_with`` both take ``jax.grad`` of a real
pairing, and their docstrings say that in a real domain that gradient and
the VJP pullback are the same map, with no conjugate-transpose convention
to disagree about. JAX returns the CONJUGATE gradient for a complex input,
so that sentence would have quietly become false the day a complex latent
arrived. Splitting first means the gradient is always taken over real
leaves and the sentence stays true -- the alternative was to keep the
solve in C and carry a conjugation convention through every one of them.
The treedef is what every ``jax.tree.map`` in the solver aligns against,
so exact invertibility is the one property that must hold:
``join(split(x)) is x`` element-wise, pinned by
``tests/exact/test_complex.py``.
Returns:
``(split, join)``. ``split`` maps ``{name: array}`` (complex where
declared) to ``{name: array | (re, im)}``; ``join`` inverts it.domain_zero
Defined in bayesmith.exact.block, line 159.
domain_zero(block: LinearBlock) -> dict[str, jax.Array]A zero of the block's real degrees of freedom. In PARTS space, not the domain: this seeds the solver's working vector and the power iteration, both of which live where :func:`real_parts` puts them. For an all-real block the two spaces are the same thing, which is why every caller predating complex support reads unchanged.
domain_centre
Defined in bayesmith.exact.block, line 180.
domain_centre(block: LinearBlock) -> dict[str, jax.Array]The prior's centre, laid out over the domain. A zero-mean prior is wrong for most physical quantities -- a noise-wave temperature sits near 250 K, not near zero -- and shifting the prior is not the same act as shifting the model even though the two give the same Gaussian. The graph states which one was meant, so this reads it.
variance_parts
Defined in bayesmith.exact.block, line 191.
variance_parts(block: LinearBlock) -> dict[str, jax.Array]``S``'s diagonal, laid out over the domain. Block-diagonal assembly by **placement** rather than concatenation: each member's variance lands on the leaf its own parameters live on, so ``x / variance`` inside a ``jax.tree.map`` IS ``S^-1 x`` with no indices to get wrong. Laid out over the REAL degrees of freedom, so it aligns with :func:`domain_zero` and with the solver's working vector. A complex member's variance is therefore **duplicated across its real and imaginary parts**, which is what ``prior_std`` means for one: each part carries ``prior_std**2``, so the complex latent has total prior variance ``2 * prior_std**2``. That is the convention rheplicant's ``gcr_sample`` documents, and the two must agree digit for digit or the seam reports a factor of sqrt(2) as a physics result.
largest_variance
Defined in bayesmith.exact.block, line 216.
largest_variance(prior_variance: dict[str, jax.Array]) -> jax.ArrayThe biggest prior variance anywhere in the block. ``1/largest`` floors ``lambda_min`` of the normal operator: ``A^T N^-1 A`` is positive semi-definite, so the LOOSEST prior in the block is what bounds the operator from below. Taking the tightest instead would floor the estimate above the true ``lambda_min`` and report a condition number smaller than the real one -- an over-confident guard, which is the direction that costs something.
isolate
Defined in bayesmith.exact.block, line 439.
isolate(graph: Graph, names: tuple[str, ...], at: dict[str, Any]) -> Callable[[dict[str, Any]], dict[str, jax.Array]]``g: {name: x} -> {observed: loc}`` -- the block's action on the prediction.
Built on ``evaluate``, so there is exactly one forward scan in this package
and the block cannot diverge from what ``log_joint`` and the NumPyro bridge
read.unchecked_operator
Defined in bayesmith.exact.block, line 460.
unchecked_operator(graph: Graph, names: Iterable[str], at: dict[str, Any] | None=None, *, probe_gaussian: bool=True, nuts_latents: Iterable[str] | None=None) -> LinearBlockExport ``A``, ``A^T``, the offset, the data and the prior -- **unchecked**.
Args:
graph: the model.
names: the latents in the block. Solving a group JOINTLY is not the
same as alternating over its members: two latents the data barely
tells apart are resolved in one CG here, where alternation
converges at the rate of their correlation while reporting a
converged residual and a condition number of ~1 at every step.
at: values for the latents OUTSIDE the block. A block is affine
*given* them, so this fixes where it is built -- which is what
makes a Gibbs sweep possible: rebuild here every sweep at the
current values.
nuts_latents: the real global NUTS block from the completed
classification or factor plan. Required when ``graph`` carries
evidence terms: another exact block is not part of NUTS merely
because it lies outside this block.
probe_gaussian: run :func:`~bayesmith.exact.gaussian.check_gaussian`
on every block member and every observed node. Default ``True``,
which is what every P3a call site already gets.
**Pass ``False`` only inside a traced Gibbs sweep, and only after
the same check has run once at compile time.** The probe evaluates
``log_prob`` at concrete offsets and raises on a mismatch, so it
does ``bool(jnp.all(...))`` -- a ``TracerBoolConversionError``
under ``jit``, not a slow path. There are TWO call sites (per
member in ``_env_before``, per observed node here) and disabling
only one still dies.
What each of this module's three entry points checks:
================================== ========== ===========
entry point linearity gaussianity
================================== ========== ===========
``linear_operator`` yes yes
``unchecked_operator`` no yes
``unchecked_operator(..., False)`` no no
================================== ========== ===========
Raises:
GraphError: if ``names`` is empty, repeats a latent, or names
something that is not a latent; if ``at`` names something that
is not a latent, names a member of this block (silently discarded
otherwise -- see :func:`_validated_at`), or omits a latent that
is outside the block; or if the graph has no observed node, so
there is nothing for a linear block to condition on; or if an
evidence term covers a non-NUTS latent, whose exact or
conditional solve would omit that term; or if an evidence graph
is supplied without the global ``nuts_latents`` witness.
NotGaussian: if a member or an observed node is not a diagonal
Gaussian, or if a member is an ancestor of another member.
StructureError: if a node's own ``log_prob`` disagrees with the
``loc``/``scale`` read off it.
Note:
The name is the warning. This does **not** check the ``linear_in``
declaration, so it will happily export an operator for a block that is
not affine at all and hand back a confident, wrong posterior.
:func:`bayesmith.exact.linearity.linear_operator` is the entry point
that checks first, and is what callers should reach for. This one is
for inside a Gibbs sweep, where the check has been hoisted out of the
loop deliberately.
Trace-safety here is a claim about THIS library's code, not about the
graph it is handed. A ``Deterministic`` whose ``fn`` does
``float(x) > 0`` raises ``ConcretizationTypeError`` under ``jit`` no
matter what this keyword says -- and a dispatcher runs arbitrary user
``fn`` under trace.