Skip to main content

PreparedObservation

observations.PreparedObservation(
    params,
    init,
    update,
    output,
    alignment_steps=None,
)

Define a temporal monitor’s block execution recipe.

prepare_observation returns this immutable recipe. The native solver and apply_observation use the same callbacks, so scientific update code is implemented once.

Fixed preparation values such as indices, strides, kernels, and prepared child recipes are captured by the callbacks. Current live values arrive through params.

Parameters

Name Type Description Default
params Any Initial live parameter PyTree. Native preparation publishes it as config.monitor so compatible values can be swept or differentiated. required
init Callable JAX-compatible callback init(params) -> state. It runs once at each model invocation, before block scanning, and constructs the scientific state from the current live parameters. Stateless monitors return None. required
update Callable JAX-compatible callback update(state, block, params) -> (next_state, samples). block is shaped [input_samples, channels, nodes] and samples must be rank three with the declared completed-period sample count. The callback must preserve the state signature and node axis. required
output ObservationOutput | dict[str, ObservationOutput] Metadata for the emitted series. A mapping is reserved for prepared composites such as JointObservation; an ordinary monitor emits one ObservationOutput. required
alignment_steps int | None Optional positive alignment in input samples. It must be a multiple of every output cadence. Regular solver blocks and truncation windows must honor this alignment; only the final incomplete tail may be shorter. None

Notes

Preparation and abstract validation may trace both callbacks more than once. Callbacks execute inside JAX transforms and must avoid host transfers, mutation, and Python control flow that depends on traced values. State and output structures and shapes are static. init may choose state dtypes from live parameters, after which every update must preserve that invocation’s state dtypes.