pyhgf.utils.vectorised_belief_propagation.batch_step#

pyhgf.utils.vectorised_belief_propagation.batch_step(network, opt_state, x, y, optimiser=None, learning_kind='precision_weighted', update_precisions=True, time_step=1.0, predicted=None, sample_weight=None, synaptic_uncertainty_settings=None, weight_reuse=1.0)#

One batch-synchronous learning step over many samples at once.

Every sample in the batch is processed from the same state template — same weights, same precisions — through the same sweeps as sample_step(), under jax.vmap, so samples are exchangeable and nothing depends on their order. The per-sample results are then averaged and applied once, so the batch counts as a single observation:

  • the mean weight gradient drives one optimiser step (skipped when optimiser is None);

  • the mean precision increments are added to the carried fields (skipped when update_precisions is False, e.g. to keep the carried precisions pinned when comparing against backpropagation).

Averaging (rather than summing) makes the result invariant to repeating the batch: the same samples twice produce the same step.

Parameters:

  • network (VectorisedNetwork) – The state template shared by every sample in the batch.

  • opt_state (Optional[optax.OptState]) – The optimiser state, or None when optimiser is None.

  • x (jnp.ndarray) – Predictors, shape (batch, n_input_features).

  • y (jnp.ndarray) – Observations, shape (batch, n_output_features).

  • optimiser (Optional[optax.GradientTransformation]) – Optax optimiser for the weight step. None freezes the weights.

  • learning_kind (str) – Weight-gradient mode.

  • update_precisions (bool) – Whether to carry the batch-averaged precision increments into the returned network.

  • time_step (float) – Inference time step, applied once per batch.

  • sample_weight (Optional[jnp.ndarray]) –

    Optional per-sample weights, shape (batch,). The batch mean becomes a weighted mean whose denominator is sample_weight.sum() rather than the row count, so rows that carry no information do not dilute the update.

    This exists because “average over the batch” is ambiguous once a caller pads. A padded row contributes a zero gradient either way, but with a plain mean it still counts in the denominator, so the effective step shrinks by the padding fraction. Anything that hands this function a variable-length batch — a token sequence, a masked objective, a ragged observation — is affected, and the symptom is a systematic gradient scale error rather than noise. Pass the mask as weights to make the reduction mean-over-contributing-rows instead.

    None (default) keeps the plain mean, so existing behaviour is unchanged.

  • predicted (Optional[tuple]) – Optional per-sample predicted states from batched_prediction_states() (one batched LayerState per element). When given, the internal prediction sweep is skipped and the update starts from these states — the forward pass a caller has already run is not repeated. x is ignored in that case.

  • synaptic_uncertainty_settings (Optional[SynapticUncertaintySettings]) – When given, the weight-belief rule runs in place of optimiser: each element’s mean and accumulated precision advance together and the optimiser state is left untouched (see pyhgf.updates.vectorised.learning.resolve_synaptic_uncertainty_settings() ). learning_kind still selects the gradient the rule descends.

  • weight_reuse (float) –

    How many times each weight matrix is applied per sample, default 1.0 (once, the ordinary case).

    This exists because “average over the batch” is also ambiguous when one weight matrix is reused several times per sample. A weight shared across k positions of a sample sees k rows per sample, so the plain mean divides by k more than that weight’s true per-sample quantities, which sum over its k uses and average only over samples. Pass k to recover those sums. The caller owns the count, since only it knows how the rows were built (see pyhgf.model.conv.conv_block(), which passes the patch count).

    Both halves of the step are rescaled, so the weight-belief rule of synaptic_uncertainty stays internally consistent: the gradient because the chain rule sums a shared weight’s uses, and the importance because the curvature those uses impose accumulates the same way. Rescaling only the gradient would move the mean k times faster while the belief tightened at the one-use rate, leaving a step k times too large once accumulated curvature dominates the prior.

    The importance half carries a modelling assumption the gradient half does not. Summing the gradient over uses is the chain rule; summing curvature over them treats the k uses as independent observations, which overlapping convolution patches are not. Where that matters, the same correction can be had with a smaller k.

Returns:

  • network – The template advanced by one batch: new weights and, if requested, new precisions. Everything else is untouched (it is rewritten by the sweeps on the next call anyway).

  • opt_state – The advanced optimiser state (None if no optimiser was given).

  • input_errors – Per-sample prediction errors at the input layer, shape (batch, n_input_features) — the messages a caller passes to whatever sits behind this network.

Return type:

tuple[VectorisedNetwork, Optional[optax.OptState], jnp.ndarray]