bayesmith.dispatch.amortized

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

Module contract

Encode a trained NeuralPosterior as an ESTIMATOR artifact, and back.

R2 Task 7, and the one place R2 reaches into R1's frozen schema.  A trained
:class:`bayesmith.amortize.NeuralPosterior` is an eqx.Module plus a static
embed callable plus standardization arrays, and none of that has a canonical
artifact form -- a callable has no canonical form (§0 ruling 4), and an
eqx.Module is a runtime object.  It is therefore encoded as the estimator arm
of :class:`bayesmith.artifacts.results.FittedConditionalPosterior`: a
REFERENCE to an artifact of the new :class:`ArtifactKind.ESTIMATOR`, never the
object itself (§0.7).

The estimator artifact is a frozen dataclass holding an opaque
`blob` -- the bytes `equinox.tree_serialise_leaves` writes for the MLP's
parameter leaves -- and a `manifest` -- the static facts a deserializer needs
to rebuild the estimator: the embed callable's importable identity, the mixture
shape, the minimum scale, the MLP's own structure, and the standardization
arrays taken from the training bank.  The blob is `bytes` to the codec, so
the artifacts layer never imports JAX or equinox; this module is where that
stack is allowed, on the dispatch side of the bridge.

Layering: this module imports equinox/JAX (the bridge), reads
`bayesmith.amortize`, and reaches the artifact protocol
(`bayesmith.artifacts.*`).  It is NOT imported by `bayesmith.artifacts`, and
nothing in the artifacts layer sees a single jax array.

EstimatorManifest EstimatorArtifact encode_estimator decode_estimator estimator_ref fitted_conditional_posterior

EstimatorManifest

Defined in bayesmith.dispatch.amortized, line 57.

class EstimatorManifest
embed_identity: str
n_components: int
n_params: int
min_scale: float
in_size: int
out_size: int
width_size: int
depth: int
theta_mean: np.ndarray
theta_scale: np.ndarray
data_mean: np.ndarray
data_scale: np.ndarray

The static half of an estimator artifact: what rebuilds the estimator.

Everything here is canonical (str / int / float / numpy arrays), so the
whole manifest round-trips through the codec without the estimator object.
The embed callable is recorded as its importable identity -- module and
qualname -- never as the callable itself, and resolved back on deserialize.

EstimatorArtifact

Defined in bayesmith.dispatch.amortized, line 104.

class EstimatorArtifact
artifact_id: str
revision: int
blob: bytes
manifest: EstimatorManifest

A fitted estimator, as identity + an opaque blob + a static manifest.

The `blob` is whatever `equinox.tree_serialise_leaves` wrote for the
MLP's parameter leaves -- opaque bytes to the artifact codec, decoded only
by :func:`decode_estimator`.  `artifact_id` and `revision` give it the
same version identity every artifact carries, so :func:`estimator_ref`
can point at it.

encode_estimator

Defined in bayesmith.dispatch.amortized, line 173.

encode_estimator(posterior: NeuralPosterior, *, artifact_id: str | None=None) -> EstimatorArtifact

Serialize a trained posterior's estimator into an ESTIMATOR artifact.

The MLP's parameter leaves are written with `equinox.tree_serialise_leaves`
into an opaque blob; every static fact (embed identity, mixture shape,
minimum scale, MLP structure, standardization) goes into the manifest.  The
result is a frozen dataclass the artifact codec can dump and load, with no
jax or equinox object in it.

decode_estimator

Defined in bayesmith.dispatch.amortized, line 210.

decode_estimator(artifact: EstimatorArtifact) -> NeuralPosterior

Rebuild the NeuralPosterior an estimator artifact encodes.

The manifest's facts rebuild the MLP's structure and the standardization;
`equinox.tree_deserialise_leaves` fills the weights back into a fresh MLP
of the same shape; the embed callable is re-imported from its identity.  The
returned estimator samples exactly as the one that was encoded.

estimator_ref

Defined in bayesmith.dispatch.amortized, line 243.

estimator_ref(artifact: EstimatorArtifact) -> ArtifactRef

The version reference to an estimator artifact: id AND revision (§0.2).

fitted_conditional_posterior

Defined in bayesmith.dispatch.amortized, line 254.

fitted_conditional_posterior(posterior: NeuralPosterior, *, simulation_bank_ref: ArtifactRef, training_run_id: str, validation_report_refs: tuple[ArtifactRef, ...]=(), artifact_id: str | None=None) -> tuple[FittedConditionalPosterior, EstimatorArtifact]

Encode a trained posterior as a FittedConditionalPosterior reference.

Returns the fitted reference AND the estimator artifact it points at, so the
caller can persist the latter and keep the former's `estimator_ref` in
agreement.  The simulation bank, training run and validation reports are
references the caller already holds; only the estimator is produced here,
because that is the one thing §0.7 adds to the frozen schema.