pyhgf.utils.vectorised_belief_propagation.propagation_step#

pyhgf.utils.vectorised_belief_propagation.propagation_step(network, opt_state, inputs, *, optimiser, time_step=1.0, learning_kind='precision_weighted', weight_update=True, synaptic_uncertainty_settings=None)[source]#

Single propagation step through the network.

Belief-propagation sweep — top-down prediction, leaf prediction error, then the interleaved posterior update + prediction error bottom-up — followed by an optional weight-learning phase. Each step dispatches per element:

  • Layer: standard per-layer kernel call (unrolled).

  • LayerStack: jax.lax.scan over the stack’s slices.

Top and bottom elements must be Layer``s. A ``LayerStack’s child below (and parent above) can themselves be Layer or LayerStack; the stack-stack case requires the boundary widths to match.

Parameters:

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

  • opt_state (optax.OptState) – The current optax optimiser state.

  • inputs (tuple) – A tuple (x, y) with the predictors set on the top element and the observations clamped on the bottom element.

  • optimiser (Optional[optax.GradientTransformation]) – The optax optimiser used for the weight-learning phase.

  • time_step (float) – The time elapsed since the previous step.

  • learning_kind (str) – The weight-gradient mode passed to pyhgf.updates.vectorised.learning.learning_weights_vectorised().

  • weight_update (bool) – Whether to apply the weight-learning phase after belief propagation.

  • synaptic_uncertainty_settings (Optional[SynapticUncertaintySettings])

Returns:

A tuple ((network, opt_state), output_pred) where network and opt_state are updated and output_pred is the bottom element’s expected_mean — the prediction of the observations for this step.

Return type:

carry