bayesmith.dispatch.collapse
Source-derived signatures and contracts. See the reference index for coverage and conventions.
Module contract
Collapsed-target sampling: integrate the exact block out of the NUTS target (P6).
**Scope, narrowed before it is widened.** Collapse is a small-block specialty,
not a general accelerator. Measured in this checkout, float64, mu = th . (B x):
* at k=8 collapse already wins for c >= ~0.95; at k=512 it needs c >= ~0.9995;
* the truncated trace-log does NOT rescue large k: it is 25-33% slower than
QR by scalar-order count and 95-130% slower by CERTIFIED order (gradient
error 2.803e-06). The reason is not numerical: rho(X) reads 0.85-0.93 on
every fixture, so the order m has to be 135-212, and the hundreds of
operator applications spend everything that never materialising the matrix
saved. The root cause is that P = A S A^T is simply NOT small next to
Lambda = N, and there is no physically grounded generic Lambda selector
(D83).
So the collapsed route materialises the dense design and does one QR per
gradient -- which is exactly the side of the package exact/solve.py exists to
avoid -- and it is only worth it when the coupling is strong enough that the
Gibbs sweep's tau(c) amplification outweighs that cost. Nothing here decides
when; the cost scoreboard (dispatch/costs.py) reports the numbers, and the
caller chooses.
The arithmetic was verified end to end before this module existed:
unchecked_operator(probe_gaussian=False) -> dense_operator -> compress
-> SqrtInfo.combine(nuisance_prior) -> marginalise_arrays
against a dense slogdet oracle gives 1.8e-14 on the value and 2.6e-08 on the
gradient, and traces cleanly under jit and hessian. The
-sum(log pivots[:n_block]) folded into the offset by marginalise_arrays is
-0.5 * logdet(F_bb) -- measured -2.1914870216494005 against -2.1914870216494.
That sentence read `IS 0.5 * logdet(F_bb)` until R4, with the sign the wrong
way round. The code was right and only the prose was wrong, which is the worse
of the two: this is the sentence an evidence assembler copies, and the copy
would have been off by logdet(F_bb) with every quantity in it finite and
plausible. A docstring cannot be pinned, so the convention now has a test --
`test_the_block_determinant_enters_with_a_minus_one_half`, which asserts it
through the determinant lemma using only public API.CollapsedEvidence collapse_graph marginal_log_density observed_descendants pivots_are_finite pivots_constrain_block
CollapsedEvidence
Defined in bayesmith.dispatch.collapse, line 235.
class CollapsedEvidence(eqx.Module)
graph: Graph
exact_names: tuple[str, ...] = eqx.field(static=True)
over: tuple[str, ...] = eqx.field(static=True)The marginal likelihood log p(d | theta) as a graph evidence term. Carries the ORIGINAL graph and the integrated block, because the reduced graph it is attached to no longer contains them; log_density rebuilds the block operator from the stored graph at the passed values.
CollapsedEvidence.log_density
log_density(self, graph: Graph, values: dict[str, Any]) -> jax.Arraycollapse_graph
Defined in bayesmith.dispatch.collapse, line 252.
collapse_graph(graph: Graph, exact_names: Iterable[str], nuts_names: Iterable[str]) -> ReducedGraphReduce the graph: exact block integrated out, marginal term attached. Atomic by construction -- reduce_with_evidence returns the reduced graph and the evidence term as one value, so the data cannot be evaluated twice. nuts_names is the explicit NUTS witness the reduced graph is sampled over.
marginal_log_density
Defined in bayesmith.dispatch.collapse, line 116.
marginal_log_density(graph: Graph, exact_names: tuple[str, ...], values: dict[str, Any]) -> jax.Arraylog p(d | theta) -- the exact block integrated out of the joint. The exact block is affine given the outside latents, so its marginal likelihood is a Gaussian integral in square-root form: whiten the dense design, append the block's prior rows, and drop the leading columns. The offset that falls out of that drop is the marginal log-density, normalisation included.
observed_descendants
Defined in bayesmith.dispatch.collapse, line 67.
observed_descendants(graph: Graph, remove: Iterable[str]) -> tuple[str, ...]The observed nodes reachable from the removed latents, transitively. These, and only these, have their likelihood moved into the marginal evidence term. An observed node independent of the removed block stays an explicit likelihood in the reduced graph; absorbing it would count its density twice.
pivots_are_finite
Defined in bayesmith.dispatch.collapse, line 91.
pivots_are_finite(pivots: jax.Array) -> jax.ArrayWhether every pivot of the marginal QR is finite (D97). A non-finite pivot means the re-triangularisation itself carried nan/inf, so the marginal log-density is not a number. Under the eager marginalise() this is the finiteness check; here it is traced and handed to eqx.error_if by marginal_log_density.
pivots_constrain_block
Defined in bayesmith.dispatch.collapse, line 102.
pivots_constrain_block(pivots: jax.Array, n_block: int) -> jax.ArrayWhether the block's pivots all sit above the relative floor (D98). The floor is sqrt(eps) times the largest pivot -- the same relative floor the eager marginalise() applies. A pivot at or below it means the block does not constrain one of its own directions, so the Gaussian integral diverges and finite arithmetic would hand back a large plausible number instead of an infinity anyone notices.