quantax.nn.RefModel#

class quantax.nn.RefModel#

Bases: Module

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

__call__(x: Array) ndarray | Array | LogArray | ScaleArray#

Usual forward pass without internal quantities.

init_internal(x: Array) PyTree#

Return initial internal values for the given configuration.

ref_forward(s: Array, s_old: Array, nflips: int, internal: PyTree, return_update: bool = False) ndarray | Array | LogArray | ScaleArray | 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.

  • nflips – The number of local updates. It’s equivalent to the number of spin flips in spin systems or the number of fermion operators in fermion systems.

  • internal – The internal quantities.

  • return_update – Whether to return the updated internal quantities.