quantax.nn.RefModel#

class quantax.nn.RefModel#

Bases: Module

The model that allows accelerated forward pass through local updates and internal quantities.

__call__(s: Array) → NDArray | Array | LogArray | ScaleArray#

Usual forward pass without internal quantities.

property use_ref: bool#

Whether to use reference implementation for local updates. Default to True.

init_internal(s: Array) → tuple[NDArray | Array | LogArray | ScaleArray, PyTree]#

Return initial wavefunction and internal values for the given configuration.

Returns:

A tuple of (initial wavefunction, internal quantities).

property required_update_modes: tuple[str, ...]#

The required update modes for accelerated ref_forward pass.

ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: Literal[False] = False) → NDArray | Array | LogArray | ScaleArray#
ref_forward(s: Array, s_old: Array, update_mode: dict[str, Any], internal: PyTree, return_update: Literal[True]) → tuple[NDArray | Array | LogArray | ScaleArray, PyTree]

Accelerated forward pass through local updates and internal quantities.

Parameters:

  • s – The new configuration.

  • s_old – The old configuration.

  • update_mode – A dictionary specifying the update mode. For instance, {"nflips": 2} indicates that there are 2 local updates.

  • internal – The internal quantities.

  • return_update – Whether to return the updated internal quantities.