bayesmith.graph.trace

Source-derived signatures and contracts. See the reference index for coverage and conventions.

Module contract

Tracing a model function into an explicit graph.

The graph object is the ground truth; this module is sugar for producing one.
Tracing runs the model **once**, so the structure it records must not depend
on values -- which is exactly the static-DAG restriction the rest of the
package relies on.

Edges are declared by passing a :class:`NodeRef` returned by an earlier
primitive, not by inspecting argument names. That keeps ``parents`` explicit
and ordered, which is what ``fn`` and ``dist_fn`` are called with. A
``NodeRef`` (and likewise a ``PlateRef``) is bound to the recorder of the
``trace(...)`` call that created it, so a handle from a different trace is
refused rather than silently resolving by name against whatever the active
trace happens to have declared -- see ``_parent_names`` and ``_plate_names``.

**Not thread-safe.** ``_STACK`` is a single process-global list. It supports
nested ``trace(...)`` calls on one thread (an inner trace pushes and pops its
own recorder around the outer one), but concurrent ``trace(...)`` calls from
multiple threads race on the same stack and can attribute one thread's nodes
to another thread's graph. Confine tracing to a single thread at a time.

NodeRef PlateRef plate const det sample observe joint_prior trace

NodeRef

Defined in bayesmith.graph.trace, line 35.

NodeRef(name: str, owner: _Recorder) -> None

A handle to a declared node. Pass one to declare an edge.

Bound to the recorder of the ``trace(...)`` call that created it. A
parent must be a ``NodeRef`` minted by the *currently active* trace --
enforced in ``_parent_names`` -- so a handle leaked in from an outer or
an earlier trace is refused instead of silently resolving against a
same-named node the active trace happens to declare.

PlateRef

Defined in bayesmith.graph.trace, line 55.

PlateRef(name: str, owner: _Recorder) -> None

A handle to a declared plate. Pass one as ``plate=`` to place a node.

Owner-checked the same way as :class:`NodeRef`, for the same reason:
without it, a plate name that merely *coincides* with one declared by a
different (or an outer) trace would bind silently -- enforced in
``_plate_names``.

plate

Defined in bayesmith.graph.trace, line 160.

plate(name: str, size: int) -> PlateRef

Declare a repetition axis and return a handle to it.

const

Defined in bayesmith.graph.trace, line 170.

const(name: str, value: Any, *, plate: PlateRef | Iterable[PlateRef] | None=None) -> NodeRef

Declare a fixed input node.

det

Defined in bayesmith.graph.trace, line 185.

det(name: str, fn: Callable[..., Any], *parents: NodeRef, linear_in: Iterable[str]=(), plate: PlateRef | Iterable[PlateRef] | None=None) -> NodeRef

Declare a deterministic node; ``linear_in`` is optional legacy metadata.

The compiler discovers supported affine structure from the computation,
even when no parent is named in ``linear_in``.

sample

Defined in bayesmith.graph.trace, line 209.

sample(name: str, dist_fn: Callable[..., Any], *parents: NodeRef, support: Support | None=None, depends_on_prediction: bool=True, plate: PlateRef | Iterable[PlateRef] | None=None) -> NodeRef

Declare a latent probabilistic node.

``support`` and ``depends_on_prediction`` are dispatch-axis claims about
the model, exactly like ``det``'s ``linear_in`` -- see
:class:`~bayesmith.graph.nodes.Probabilistic` for what each default
means and why. Nothing in P1 reads either.

observe

Defined in bayesmith.graph.trace, line 238.

observe(name: str, dist_fn: Callable[..., Any], *parents: NodeRef, obs: Any, mask: Any=None, support: Support | None=None, depends_on_prediction: bool=True, plate: PlateRef | Iterable[PlateRef] | None=None) -> NodeRef

Declare a probabilistic node conditioned on data.

``support`` and ``depends_on_prediction`` are dispatch-axis claims about
the model, exactly like ``det``'s ``linear_in`` -- see
:class:`~bayesmith.graph.nodes.Probabilistic` for what each default
means and why. Nothing in P1 reads either.

``mask`` is boolean and shaped like ``obs``, ``True`` where the sample was
actually taken; see
:attr:`~bayesmith.graph.nodes.Probabilistic.observed_mask` for why an
unobserved sample is declared here and not as an infinite scale.

joint_prior

Defined in bayesmith.graph.trace, line 275.

joint_prior(prior: Any) -> None

Declare a prior over SEVERAL latents at once -- one per graph.

A latent's own ``dist_fn`` is the prior on that one quantity. A Jeffreys
prior is a different shape of object: one density over a named block, read
off the forward model and the noise. Declaring it here rather than adding
``numpyro.factor`` by hand in a model is what makes the two scans of a
graph -- :func:`~bayesmith.graph.evaluate.log_joint` and
:func:`~bayesmith.bridge.numpyro_bridge.to_numpyro` -- read one
declaration instead of one of them silently missing it.

Order does not matter: it is a term of the joint, not a node, so it may be
declared before or after the latents it covers. The covered latents must
themselves be declared FLAT; the prior refuses a covered latent carrying a
proper density where it is evaluated, since that is two priors on one
quantity.

Raises:
    GraphError: if a second one is declared. **Two Jeffreys blocks are not
        two independent factors.** Each is the CONDITIONAL prior of its
        block given the other latents, and a product of conditionals is in
        general not the joint density of anything -- it is neither
        declaration's prior, and nothing downstream reports the
        difference. One block, named explicitly, is the declaration this
        package can be held to. A model that genuinely wants both blocks
        covered declares ONE prior over their union, which is a different
        (and often non-existent) density -- see ``JeffreysPrior.over``.

trace

Defined in bayesmith.graph.trace, line 319.

trace(model_fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Graph

Run ``model_fn`` once and return the graph it declared.