mixle.inference.transaction module

Transactional snapshots for mutable model state used inside EM updates.

Most mixle distributions are immutable value objects: an M-step returns a new distribution and leaves the previous iterate untouched. Torch-backed leaves are the important exception. Their estimators update a shared nn.Module in place, so an objective gate needs to snapshot that state before proposing a step and restore it when the proposal is rejected.

The helpers here deliberately recognize the small state_dict / load_state_dict protocol instead of importing torch. This keeps torch an optional dependency and also works for compatible stateful modules.

stateful_modules(*roots)

Return unique mutable modules reachable from roots.

Traversal stops at a module boundary. Walking through a torch module’s own object graph would be both expensive and redundant because state_dict is already the module’s complete transactional state.

Parameters:

roots (Any)

Return type:

tuple[Any, …]

has_mutable_state(*roots)

Whether any torch-like mutable module is reachable from roots.

Parameters:

roots (Any)

Return type:

bool

class MutableStateSnapshot(entries)

Bases: object

Restorable state for all mutable modules reachable from a model tree.

Parameters:

entries (tuple[tuple[Any, dict[str, Any], bool | None], ...])

class AlgorithmStateSnapshot(dict_entries, slot_entries, rng_entries, preserved)

Bases: object

In-place snapshot of a strategy object graph and every reachable NumPy RNG.

Parameters:

• dict_entries (tuple[tuple[Any, dict[str, Any]], ...])

• slot_entries (tuple[tuple[Any, str, Any], ...])

• rng_entries (tuple[tuple[Any, Any], ...])

• preserved (tuple[Any, ...])

restore()

Restore captured objects without replacing their externally visible identities.

Return type:

None