bayesmith.dispatch.plan

CONVERGENCE_TARGET

Defined in bayesmith.dispatch.plan, line 85.

CONVERGENCE_TARGET: float = 0.001

KAPPA_PROBE_SIGMAS

Defined in bayesmith.dispatch.plan, line 103.

KAPPA_PROBE_SIGMAS: tuple[float, ...] = (-3.0, -1.0, 1.0, 3.0)

kappa_upper

Defined in bayesmith.dispatch.plan, line 139.

kappa_upper(kappa: float | tuple[float, float]) -> float

The end of ``kappa`` that ``tol`` is derived from.

An interval's UPPER end, never its lower one. Taking ``lo`` gives a
``tol`` that is too loose by the width of the interval -- 34x on
``indirect_ancestor`` -- and too loose is the direction that stops CG
early, returns a posterior that is too narrow, and does it inside a sweep
where the convergence guard has been hoisted out and nothing notices.

tol_for

Defined in bayesmith.dispatch.plan, line 151.

tol_for(kappa: float | tuple[float, float]) -> float

``CONVERGENCE_TARGET / kappa`` -- the one place the division lives.

working_epsilon

Defined in bayesmith.dispatch.plan, line 189.

working_epsilon(graph: Graph, names: tuple[str, ...], at: dict[str, Any]) -> float

Machine epsilon of the dtype the block's own arithmetic runs at.

Read off ``block.offset``, which is the same leaf
:func:`~bayesmith.exact.solve.wiener_solve`'s own precision guard reads
(``jnp.result_type(*jax.tree.leaves(block.offset))``). Deliberately the
same expression and not an equivalent one: the plan's verdict about
whether ``tol`` is reachable and the solver's verdict about whether
``require_convergence`` is reachable must be verdicts about one number,
or the plan can promise an accuracy the solve then refuses.

It is a property of the GRAPH rather than of this module, because
``const`` and ``observe`` capture their arrays at trace time: a graph
traced outside ``jax.enable_x64(True)`` stays float32 however it is later
compiled.

kappa_interval

Defined in bayesmith.dispatch.plan, line 261.

kappa_interval(graph: Graph, names: tuple[str, ...], *, env: dict[str, Any] | None=None, key: jax.Array | None=None) -> tuple[float, float, str]

``(lo, hi, note)`` -- how far the block's conditioning moves, and what was missed.

The centre is always evaluated, so ``lo == hi`` exactly when nothing
moved kappa: a block spanning every latent has no outside latent to sweep,
and one whose outside latents are all non-Gaussian has none this module
can sweep. Both print as a single number rather than a degenerate
interval, and the second says so in ``note``.

**A sampled interval, not a bound.** Each outside latent is moved on its
own, from the centre, along the UNIFORM direction (every element of a
plated latent displaced together) -- ``1 + len(KAPPA_PROBE_SIGMAS) *
outside`` evaluations rather than a product grid, and no contrast
direction within a plate. Every outside latent in this package's fixture
suite is scalar, so the contrast case is undocumented by measurement
rather than covered.

Args:
    graph: the model.
    names: the exact block's members.
    env: :func:`~bayesmith.dispatch.classify.prior_environment`, if it has
        already been built. Passed in rather than rebuilt so the plan
        anchors where ``partition`` classified.
    key: PRNG key for ``condition_bound``'s power iteration. Fixed by
        default, so the interval is reproducible.

Block

Defined in bayesmith.dispatch.plan, line 344.

class Block(eqx.Module)
latents: tuple[str, ...] = eqx.field(static=True)
method: str = eqx.field(static=True)
reason: str = eqx.field(static=True)
linearity: dict | None = eqx.field(static=True, default=None)
kappa: float | tuple[float, float] | None = eqx.field(static=True, default=None)
tol: float | None = eqx.field(static=True, default=None)
epsilon: float | None = eqx.field(static=True, default=None)
structure: dict | None = eqx.field(static=True, default=None)

One group of latents and the method the graph selected for it.

Attributes:
    latents: the members, sorted.
    method: ``"gcr"``, ``"gcr+snis"``, ``"gcr+mh"``, ``"log-gcr"`` or
        ``"nuts"``. FIVE values, from two producers:
        :class:`~bayesmith.dispatch.classify.Classification` chooses the
        first four, and :func:`~bayesmith.dispatch.factor.factor_partition`
        also emits ``"log-gcr"`` without going through
        :class:`Classification` at all. Neither ``_LABELS`` here nor
        ``FACTOR_METHODS`` there enumerates all five, and that is
        deliberate on both sides -- each says in its own docstring which
        row it leaves out and why -- so this attribute is the one place
        the union is written down.
    reason: why -- naming members on a refusal, plus whatever the kappa
        sweep could not reach.
    linearity: ``check_linearity``'s per-at-point errors, or ``None``.
    kappa: a single conditioning bound, or ``(lo, hi)`` when a latent
        outside the block moves it. ``None`` for a NUTS block, which is
        solved by no linear system at all.
    tol: ``CONVERGENCE_TARGET / kappa_upper(kappa)``.
    epsilon: machine epsilon of the dtype this block's arithmetic runs at,
        from :func:`working_epsilon`. ``None`` for a NUTS block, which
        solves nothing. Stored rather than re-derived at print time
        because it is a property of the GRAPH -- ``const`` and ``observe``
        capture their arrays at trace time -- and a plan is meant to be a
        record of what was measured, not a live query.

Block.tol_attainable

tol_attainable(self) -> bool

Whether CG can reach ``tol`` at all in this block's own arithmetic.

``False`` when ``kappa_upper(kappa) * epsilon >= CONVERGENCE_TARGET``:
the error a solve delivers is ``kappa`` times its relative residual,
the residual cannot go below the arithmetic's own noise, and past this
product the target is arithmetically out of reach whatever ``tol``
says. It is the same test
:func:`~bayesmith.exact.solve.wiener_solve` makes on
``require_convergence`` (``bound * epsilon > require_convergence``),
against the target ``tol`` was derived from, so the plan's verdict and
the solver's cannot disagree.

**A statement about the guarantee, not about the answer.** The bound
is an UPPER bound on kappa and a loose one -- measured at 3676x the
true conditioning on ``two_linear_latents``, and 232,000x between the
prior centre and the GLS fixed point on ``radiometer`` -- so a
``False`` here is compatible with a solve that is in fact accurate.
Ten of this package's own fixtures read ``False`` at float32 and
``tests/dispatch/test_acceptance.py`` measures two of them
(``prior_held_direction``, ``straight_line(prior_std=1e3,
sigma=1e-2)``) delivering inside ``CONVERGENCE_TARGET`` anyway. That
is exactly why this is printed rather than raised.

``True`` where there is nothing to solve, which is abstention rather
than endorsement -- the same convention
:attr:`~bayesmith.dispatch.execute.Posterior.unreliable` uses for a
missing k-hat.

InferencePlan

Defined in bayesmith.dispatch.plan, line 527.

class InferencePlan(eqx.Module)
graph: Graph
blocks: tuple[Block, ...]
sigma_needs_rebuild: bool = eqx.field(static=True, default=False)
streaming: StreamingRoute | None = eqx.field(static=True, default=None)
ladder: LadderRecord | None = eqx.field(static=True, default=None)

What :func:`compile` produces: a partition, its reasons, and how to run it.

Attributes:
    graph: the model this plan was derived from.
    blocks: the exact block first, if there is one, then the sampled one.
    sigma_needs_rebuild: whether an observed node's scale has a latent
        ancestor outside the exact block, in which case ``noise_std``
        must be recomputed every sweep rather than hoisted.
    streaming: whether the campaign fold applies to this graph, and on
        which plate -- see
        :mod:`~bayesmith.dispatch.streaming`. MEASURED here rather than
        offered as a live query, because a plan is a record of what was
        measured; ``None`` only on a plan built by hand.

InferencePlan.exact

exact(self) -> Block | None

The block an exact method applies to, or ``None``.

InferencePlan.sampled

sampled(self) -> Block | None

The block NUTS samples, or ``None`` if the graph is fully exact.

InferencePlan.guard_hoisted

guard_hoisted(self) -> bool

Whether the exact solve runs inside a sweep with its guard off.

True exactly for a MIXED plan. The guard
(:func:`~bayesmith.exact.solve.wiener_solve`'s ``require_convergence``)
is a concrete-valued check and a Gibbs sweep runs under ``jit``, so
inside a sweep it cannot run at all -- which is why ``tol`` has to
carry the discipline on its own there.

**False does not mean the guard is running**, and the printed line
says so in those words. A fully exact plan runs its solve outside any
trace, so the guard CAN run there -- but :meth:`sample` and
:meth:`estimate` leave it off and expose it as ``require_convergence=``
instead, because ``tol`` is already derived to deliver
``CONVERGENCE_TARGET`` and a guard set to that same target compares two
numbers constructed to be equal. Measured at float32: ``straight_line``
lands at ``1.078e-07 * 924.4 = 9.97e-04`` and is accepted by 0.3%,
``two_linear_latents`` at ``2.179e-07 * 5792 = 1.262e-03`` and is
refused by 26%.

InferencePlan.sample

sample(self, key: jax.Array, *, num_samples: int=2000, num_warmup: int=1000, num_chains: int=1, chain_method: str='sequential', progress_bar: bool=False, nuts_options: Mapping[str, Any] | None=None, tol: float | None=None, maxiter: int | None=None, require_convergence: float | None=None, ess_floor: float=SNIS_ESS_FLOOR, nuts_on_collapse: bool=False, collapse: bool=False, _control: Any=None) -> Posterior

Run the plan. Section 6.4's dispatch, and nothing else decides.

Five shapes, in the order this reads them: no exact block -> NUTS;
exact block plus a sampled one -> ``HMCGibbs`` with the sweep the plan
printed, or, with ``collapse=True``, the P6 collapse arm (NUTS on the
marginal target with the exact block integrated out, then one exact
``gcr_sample`` regression per retained draw); whole graph exact with a
fixed sigma -> iid GCR draws, no chain; whole graph exact with a moving
sigma -> GCR at the GLS fixed point corrected by SNIS; and that last
one again when its Kish ESS/N falls under ``ess_floor`` -> the same
weighted sample, marked ``unreliable=True`` and saying so, or NUTS
instead if ``nuts_on_collapse``. See :data:`SNIS_ESS_FLOOR` for the
measurement that decided which of those two is the default.

Args:
    key: PRNG key. Split once, so the draws and any fallback chain do
        not share a stream.
    num_samples, num_warmup, num_chains, chain_method, progress_bar,
        nuts_options: passed to whichever sampler runs -- ``HMCGibbs``
        through :func:`~bayesmith.exact.gibbs.assemble` on the mixed
        path, ``NUTS`` through
        :func:`~bayesmith.bridge.numpyro_bridge.nuts` on both bare
        ones, under the same six names in both. ``num_warmup`` is
        ignored on the iid path, there being nothing to adapt.
    tol: CG tolerance. Defaults to the plan's own
        ``exact.tol = CONVERGENCE_TARGET / kappa``; overriding it
        overrides the discipline that number carries.
    maxiter: CG iteration cap.
    require_convergence: bound on the relative ERROR of each solve.
        **Off by default, measured.** ``tol`` is derived so that a CG
        stopping at ``tol`` delivers exactly ``CONVERGENCE_TARGET``,
        so switching this on at that same target compares two numbers
        constructed to be equal and fires on rounding: at float32
        ``straight_line`` lands at ``1.078e-07 * 924.4 = 9.97e-04``
        and passes, ``two_linear_latents`` at ``2.179e-07 * 5792 =
        1.262e-03`` and does not -- two fixtures of the same shape,
        0.3% inside and 26% outside. It is also a property of the
        OPERATOR, so on the draw path it would be re-measured, at
        ``POWER_ITERATIONS`` operator applications, once per draw.
    ess_floor: the Kish ESS/N under which the SNIS path declares the
        correction collapsed.
    nuts_on_collapse: on a collapse, discard the weighted sample and
        sample the whole graph with NUTS instead. **Off by default,
        measured.** At the collapsed cell ``plated_radiometer(n=25,
        kappa=0.4)``, N=1200, against exact per-coordinate quadrature,
        the weighted answer's worst coordinate is 1.40 posterior sd
        from the truth and the NUTS replacement's is 18.5 -- while
        NUTS's chain ESS (33) exceeds the Kish ESS (14), so the
        diagnostic that fires prefers the worse answer. Substituting
        is therefore something to ask for, with the collapse already
        reported in ``unreliable`` and ``reason`` for a caller who
        wants to decide for themselves.
    collapse: run the P6 collapse arm on a MIXED plan instead of the
        Gibbs sweep -- NUTS on the marginal target with the exact block
        integrated out, then one exact gcr_sample regression per
        retained draw. **Off by default.** Only valid when the exact
        block's method is "gcr" (constant sigma); a
        prediction-dependent block refuses, because marginalising it
        would freeze sigma at the prior centre and return a plausible
        but wrong marginal. See bayesmith.dispatch.collapse.collapse_graph.

Returns:
    A :class:`Posterior`, whose ``method`` is what RAN. Its
    ``cost`` carries the P7 reconciliation ledger -- predicted
    interval, measured seconds per effective sample, and which of the
    prediction's own terms dominated it -- on a plan compiled with
    ``strategy="cost"``, and is ``None`` on every declared plan,
    which made no prediction to reconcile against.

InferencePlan.estimate

estimate(self, *, tol: float | None=None, maxiter: int | None=None, reweight_tol: float | None=None, min_reweights: int=MIN_REWEIGHTS, max_reweights: int=MAX_REWEIGHTS, require_convergence: float | None=None) -> Estimate

The point estimate, where one exists. Section 6.5's dispatch.

Whole graph exact with a fixed sigma is one
:func:`~bayesmith.exact.solve.wiener_solve`; whole graph exact with a
prediction-dependent sigma is
:func:`~bayesmith.exact.gls.iterative_gls`'s fixed point. Anything
mixed is REFUSED and pointed at :meth:`sample`.

Both paths are spelled through ``iterative_gls``, whose
``depends_on_prediction=False`` branch *is* the single ``wiener_solve``
-- so :class:`Estimate`'s five fields come from one place rather than
two, and the difference between the branches stays visible in
``iterations`` (1 against several) rather than being invisible in
which function was called.

Args:
    tol, maxiter, require_convergence: as for :meth:`sample`; ``tol``
        defaults to the plan's own and the guard is off for the same
        measured reason.
    reweight_tol, min_reweights, max_reweights: the outer
        fixed-point iteration's own settings.

Raises:
    ConvergenceError: if the reweighting did not reach a fixed point.
        P3a defined this class and left it unraised, leaving the
        promotion of ``GLSResult.converged`` to its caller; this is
        that caller.
    NotImplementedError: for a graph any part of which is only
        samplable. A point estimate there is a MAP, which is P5.

compile

Defined in bayesmith.dispatch.plan, line 878.

compile(graph: Graph, *, key: jax.Array | None=None, strategy: Literal['declared', 'cost']='declared', a: float=1.0, timing: TimingConstants | None=None, _classification: Classification | None=None) -> InferencePlan

Derive the plan for a graph: what runs, on which latents, and why.

Runs :func:`~bayesmith.dispatch.classify.partition`, then measures the
exact block's conditioning across the outside latents' own priors and
derives ``tol`` from the worst of it, and reads the block's working
precision so the plan can say whether that ``tol`` is reachable at all.
Takes no samples and forms no matrix; costs
``1 + len(KAPPA_PROBE_SIGMAS) * outside`` power iterations, plus one
further block build for :func:`working_epsilon`, on top of what
``partition`` already spends.

Then :func:`~bayesmith.dispatch.streaming.streaming_route`, which costs
one leakage probe per PLATE and nothing at all on a graph with none.
Measured on the scalar campaign: 0.278 s at four epochs and 0.342 s at
512, against this function's own 1.0-1.3 s -- flat in the campaign,
because the probe is one ``jacfwd`` of the whole plate rather than one
per epoch.

Shadows the builtin ``compile`` at module scope, which is the decided UX
-- ``bayesmith.compile(graph)`` is the name this package wants -- so
callers who need both spell it ``from bayesmith import compile as
compile_graph``.

Args:
    graph: the model.
    key: PRNG key, passed to both ``partition``'s linearity probes and
        ``condition_bound``'s power iteration. Fixed by default, so a
        plan is reproducible; and the default is
        ``condition_bound``'s own default, so a kappa printed here equals
        a kappa a caller measures directly.

Returns:
    An :class:`InferencePlan`, whose ``str`` is the readable form.