neuralqx.experimental.vqs.mc.mc_state package

class STMultiMCState(states, *, canonical_state=0, sync_model_state=False)

Bases: object

A Container for multiple head-specific MCState objects that share one parameter pytree.

This class is intentionally small and driver-oriented. It assumes you already created one MCState per head (typically from STMHHeadView models) with identical parameter tree structure, and then synchronizes their parameters so they act as different views of the same ST-MH model.

The current MultiMCState returns a list of parameter pytrees (one per state), which is exactly right for MT-MH (independent networks) but is not the ST-MH (shared trunk + heads in one joint parameter set) Ansatz. Here we expose a single parameter pytree to the base VMC driver so optimizer state and updates are computed only once.

Parameters:

  • states (list[MCState]) – List of head-specific MCState objects. They must share the same Hilbert space and the same parameter-tree structure.

  • canonical_state (int) – Index of the state whose parameters are treated as the canonical source of truth.

  • sync_model_state (bool) – If True, broadcast_from_canonical() also copies model_state from the canonical state to all other states. Keep this False unless you really need mutable model-state collections and understand the implications.

Subpackages

Submodules