pyhgf.utils.vectorised_belief_propagation.batched_prediction_states#

pyhgf.utils.vectorised_belief_propagation.batched_prediction_states(network, x)[source]#

Batched forward sweep returning the per-sample swept states.

Like batched_prediction_pass(), but keeps what the sweep computed: one batched LayerState per element (each field with a leading batch axis). Passing these to batch_step() as predicted lets the learning step start directly from them instead of repeating the forward sweep — the weights and static fields are not duplicated per sample, only the layer states are.

The states are the only output: the per-sample predictions are read from the bottom element’s expected_mean after the call. Returning that array alongside the states from the same compiled function produces incorrect values under the vmap-of-jit composition on CPU, so callers must read it from the returned states.

Parameters:

  • network (VectorisedNetwork) – The current vectorised network state.

  • x (Array) – Predictors, shape (batch, n_input_features).

Returns:

One batched LayerState per element, ordered as network.layers.

Return type:

states