jexplore.steps.modsel.swap

jexplore.steps.modsel.swap#

This module defines the class for Model Selection steps swapping chains.

Attributes#

TSwapMS

Temperature swap for model selection

Classes#

OrderByModel

Swap chains to restore model ordering (temp by temp)

Module Contents#

TSwapMS#

Temperature swap for model selection

class OrderByModel[Tepoch: jexplore.sampling.EpochMS, Tstate: jexplore.sampling.StateMS, Tsampling: jexplore.sampling.SamplingMS][source]#

Bases: jexplore.steps.step.Step[Tepoch, Tstate, Tsampling]

Swap chains to restore model ordering (temp by temp)

step(key, state)[source]#

Pseudo drawing step

Parameters:

  • key (jax.Array) – PRNG key

  • state (Tstate) – current state

Returns:

new state and the boolean mask of the chains modified by the step.

Return type:

tuple[Tstate, jax.Array]