compiled-knowledge 4.0.0a20__cp313-cp313-win32.whl

ck/sampling/marginals_direct_sampler.py

@@ -0,0 +1,113 @@
+from typing import Collection, Iterator, Dict, Sequence
+
+import numpy as np
+
+from ck.pgm import Instance
+from ck.probability.probability_space import dtype_for_state_indexes
+from ck.program.program_buffer import ProgramBuffer
+from ck.program.raw_program import RawProgram
+from ck.sampling.sampler import Sampler
+from ck.sampling.sampler_support import SampleRV, YieldF, SamplerInfo
+from ck.utils.np_extras import NDArray, NDArrayNumeric
+from ck.utils.random_extras import Random
+
+
+class MarginalsDirectSampler(Sampler):
+
+    def __init__(
+            self,
+            sampler_info: SamplerInfo,
+            raw_program: RawProgram,
+            rand: Random,
+            rv_idx_to_result_offset: Dict[int, int],
+    ):
+        super().__init__(sampler_info.rvs, sampler_info.condition)
+        self._yield_f: YieldF = sampler_info.yield_f
+        self._rand: Random = rand
+        self._program_buffer = ProgramBuffer(raw_program)
+        self._sample_rvs: Sequence[SampleRV] = tuple(sampler_info.sample_rvs)
+        self._chain_rvs: Sequence[SampleRV] = tuple(
+            sample_rv for sample_rv in sampler_info.sample_rvs if sample_rv.copy_index is not None)
+        self._state_dtype = dtype_for_state_indexes(self.rvs)
+        self._max_number_of_states: int = max((len(rv) for rv in self.rvs), default=0)
+        self._slots_1: Collection[int] = sampler_info.slots_1
+
+        self._marginals: Sequence[NDArrayNumeric] = tuple(
+            self._program_buffer.results[
+            rv_idx_to_result_offset[sample_rv.rv.idx]
+            :
+            rv_idx_to_result_offset[sample_rv.rv.idx] + len(sample_rv.rv)
+            ]
+            for sample_rv in sampler_info.sample_rvs
+        )
+        # Set up the input slots to 0 or 1 to respect conditioning and initial Markov chain states.
+        slots: NDArray = self._program_buffer.vars
+        for slot in sampler_info.slots_0:
+            slots[slot] = 0
+        for slot in sampler_info.slots_1:
+            slots[slot] = 1
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        yield_f = self._yield_f
+        rand = self._rand
+        sample_rvs = self._sample_rvs
+        chain_rvs = self._chain_rvs
+        program_buffer = self._program_buffer
+        slots: NDArray = program_buffer.vars
+        marginals = self._marginals
+        slots_1 = self._slots_1
+
+        # Set up working memory buffer
+        states = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+
+        def compute() -> float:
+            # Compute the program results based on the current input slot values.
+            # Return the WMC.
+            return program_buffer.compute().item(-1)
+
+        while True:
+            wmc: float = compute()
+            rnd: float = rand.random() * wmc
+
+            for sample_rv in sample_rvs:
+                index: int = sample_rv.index
+                if index > 0:
+                    # No need to execute the program on the first time through
+                    # as it was done just before entering the loop.
+                    wmc = compute()
+
+                rv_dist: NDArray = marginals[sample_rv.index]
+
+                rv_dist_sum: float = rv_dist.sum()
+                if rv_dist_sum <= 0:
+                    raise RuntimeError('zero probability')
+                rv_dist *= wmc / rv_dist_sum
+
+                state_index: int = -1
+                for i in range(len(sample_rv.rv)):
+                    w = rv_dist.item(i)
+                    if rnd < w:
+                        state_index = i
+                        break
+                    rnd -= w
+                assert state_index >= 0
+
+                for slot in sample_rv.slots:
+                    slots[slot] = 0
+                slots[sample_rv.slots[state_index]] = 1
+                states[index] = state_index
+
+            yield yield_f(states)
+
+            # Reset the one slots for the next iteration.
+            for slot in slots_1:
+                slots[slot] = 1
+
+            # Copy chain pairs for next iteration.
+            # (This writes over any initial chain conditions from slots_1.)
+            for sample_rv in chain_rvs:
+                rv_slots = sample_rv.slots
+                prev_state_idx: int = states.item(sample_rv.copy_index)
+                for slot in rv_slots:
+                    slots[slot] = 0
+                slots[rv_slots[prev_state_idx]] = 1

ck/sampling/sampler.py

@@ -0,0 +1,62 @@
+from abc import ABC, abstractmethod
+from itertools import islice
+from typing import Sequence, Iterator
+
+from ck.pgm import RandomVariable, Instance, Indicator
+
+
+class Sampler(ABC):
+    """
+    A Sampler provides an unlimited series of samples for one or more random variables.
+    The random variables being sampled are provided as a tuple via the `rvs` property.
+
+    A Sampler will either iterate over Instance objects, where each instance is co-indexed
+    with `self.rvs`, or may iterate over single state indexes. Whether a Sampler iterates
+    over Instance objects or single state indexes is determined by the implementation.
+    If iterating over single state indexes, then `len(self.rvs) == 1`.
+    """
+    __slots__ = ('_rvs', '_condition')
+
+    def __init__(self, rvs: Sequence[RandomVariable], condition: Sequence[Indicator]):
+        """
+        Args:
+            rvs: a collection of the random variables being
+                sampled, co-indexed with each sample provided by `iter(self)`.
+            condition: condition on `rvs` that are compiled into the sampler.
+        """
+        self._rvs: Sequence[RandomVariable] = tuple(rvs)
+        self._condition: Sequence[Indicator] = tuple(condition)
+
+    @property
+    def rvs(self) -> Sequence[RandomVariable]:
+        """
+        What random variables are being sampled.
+
+        Returns:
+            the random variables being sampled, co-indexed with each sample from `iter(self)`.
+        """
+        return self._rvs
+
+    @property
+    def condition(self) -> Sequence[Indicator]:
+        """
+        Condition on `self.rvs` that are compiled into the sampler.
+        """
+        return self._condition
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        """
+        An unlimited series of samples from a random process.
+        Each sample is co-indexed with the random variables provided by `self.rvs`.
+        """
+        ...
+
+    def take(self, number_of_samples: int) -> Iterator[Instance] | Iterator[int]:
+        """
+        Take a limited number of samples from `iter(self)`.
+
+        Args:
+            number_of_samples: a limit on the number of samples to provide.
+        """
+        return islice(self, number_of_samples)

ck/sampling/sampler_support.py

@@ -0,0 +1,232 @@
+from dataclasses import dataclass
+from itertools import count
+from typing import Callable, Sequence, Optional, Set, Tuple, Dict, Collection
+
+from ck.pgm import Instance, RandomVariable, Indicator
+from ck.pgm_circuit.program_with_slotmap import ProgramWithSlotmap
+from ck.pgm_circuit.slot_map import SlotMap
+from ck.probability.probability_space import Condition, check_condition
+from ck.utils.map_set import MapSet
+from ck.utils.np_extras import NDArrayStates, NDArrayNumeric
+from ck.utils.random_extras import Random
+
+# Type of a yield function. Support for a sampler.
+# A yield function may be used to implement a sampler's iterator, thus
+# it provides an Instance or single state index.
+YieldF = Callable[[NDArrayStates], int] | Callable[[NDArrayStates], Instance]
+
+
+@dataclass
+class SampleRV:
+    """
+    Support for a sampler.
+    A SampleRV structure keeps track of information for one sampled random variable.
+    """
+    index: int  # index into the sequence of sample rvs.
+    rv: RandomVariable  # the random variable being sampled.
+    slots: Sequence[int]  # program input slots for indicators of the random variable (co-indexed with rv.states).
+    copy_index: Optional[int]  # for Markov chains, which previous sample rv should be copied?
+
+
+@dataclass
+class SamplerInfo:
+    """
+    Support for a sampler.
+    A SamplerInfo structure keeps track of standard information when a sampler uses a Program.
+    """
+    sample_rvs: Sequence[SampleRV]
+    condition: Sequence[Indicator]
+    yield_f: YieldF
+    slots_0: Set[int]
+    slots_1: Set[int]
+
+    @property
+    def rvs(self) -> Tuple[RandomVariable, ...]:
+        """
+        Extract the RandomVariable objects from `self.sample_rvs`.
+        """
+        return tuple(sample_rv.rv for sample_rv in self.sample_rvs)
+
+
+def get_sampler_info(
+        program_with_slotmap: ProgramWithSlotmap,
+        rvs: Optional[RandomVariable | Sequence[RandomVariable]],
+        condition: Condition,
+        chain_pairs: Sequence[Tuple[RandomVariable, RandomVariable]] = (),
+        initial_chain_condition: Condition = (),
+) -> SamplerInfo:
+    """
+    Helper for samplers.
+
+    Determines:
+    (1) the slots for sampling rvs,
+    (2) Markov chaining rvs,
+    (3) the function to use for yielding an Instance or state index.
+
+    If parameter `rvs` is a RandomVariable, then the yield function will
+    provide a state index. If parameter `rvs` is a Sequence, then the
+    yield function will provide an Instance.
+
+    Args:
+        program_with_slotmap: the program and slotmap being referenced.
+        rvs: the random variables to sample. It may be either a sequence of
+            random variables, or a single random variable.
+        condition: is a collection of zero or more conditioning indicators.
+        chain_pairs: is a collection of pairs of random variables, each random variable
+            must be in the given rvs. Given a pair (from_rv, to_rv) the state of from_rv is used
+            as a condition for to_rv prior to generating a sample.
+        initial_chain_condition: are condition indicators (just like condition)
+            for the initialisation of the 'to_rv' random variables mentioned in chain_pairs.
+
+    Raises:
+        ValueError: if preconditions of `program_with_slotmap` are incompatible with the given condition.
+
+    Returns:
+        a SamplerInfo structure.
+    """
+    if rvs is None:
+        rvs = program_with_slotmap.rvs
+    if isinstance(rvs, RandomVariable):
+        # a single rv
+        rvs = (rvs,)
+        yield_f = lambda x: x.item()
+    else:
+        # a sequence of rvs
+        rvs = tuple(rvs)
+        yield_f = lambda x: x.tolist()
+
+    # Group condition indicators by `rv_idx`.
+    conditioned_rvs: MapSet[int, Indicator] = MapSet()
+    for ind in check_condition(condition):
+        conditioned_rvs.add(ind.rv_idx, ind)
+    del condition
+
+    # Group precondition indicators by `rv_idx`.
+    preconditioned_rvs: MapSet[int, Indicator] = MapSet()
+    for ind in program_with_slotmap.precondition:
+        preconditioned_rvs.add(ind.rv_idx, ind)
+
+    # Rationalise conditioned_rvs with preconditioned_rvs
+    rv_idx: int
+    precondition_set: Set[Indicator]
+    for rv_idx, precondition_set in preconditioned_rvs.items():
+        condition_set = conditioned_rvs.get(rv_idx)
+        if condition_set is None:
+            # A preconditioned rv was not mentioned in the explicit conditions
+            conditioned_rvs.add_all(rv_idx, precondition_set)
+        else:
+            # A preconditioned rv was also mentioned in the explicit conditions
+            condition_set.intersection_update(precondition_set)
+            if len(condition_set) == 0:
+                rv_index: Dict[int, RandomVariable] = {rv.idx: rv for rv in rvs}
+                rv: RandomVariable = rv_index[rv_idx]
+                raise ValueError(f'conditions on rv {rv} are disjoint from preconditions')
+    del preconditioned_rvs
+
+    # Group initial chain indicators by `rv_idx`.
+    initial_chain_condition: Sequence[Indicator] = check_condition(initial_chain_condition)
+    initial_chain_conditioned_rvs: MapSet[int, Indicator] = MapSet()
+    for ind in initial_chain_condition:
+        initial_chain_conditioned_rvs.add(ind.rv_idx, ind)
+
+    # Check sample rvs are valid and without duplicates.
+    rvs_set: Set[RandomVariable] = set(rvs)
+    if not rvs_set.issubset(program_with_slotmap.rvs):
+        raise ValueError('sample random variables not available')
+    if len(rvs) != len(rvs_set):
+        raise ValueError('duplicate sample random variables requested')
+
+    # Check chain_pairs rvs are being sampled
+    if not rvs_set.issuperset(pair[0] for pair in chain_pairs):
+        raise ValueError('a random variable appears in chain_pairs but not in sample rvs')
+    if not rvs_set.issuperset(pair[1] for pair in chain_pairs):
+        raise ValueError('a random variable appears in chain_pairs but not in sample rvs')
+
+    # Check chain_pairs source and destination rvs are disjoint
+    if not {pair[0] for pair in chain_pairs}.isdisjoint(pair[1] for pair in chain_pairs):
+        raise ValueError('chain_pairs sources and destinations are not disjoint')
+
+    # Check no chain_pairs destination rv is a conditioned rv
+    if any(pair[1].idx in conditioned_rvs.keys() for pair in chain_pairs):
+        raise ValueError('a chain_pairs destination is conditioned')
+
+    # Check chain initial conditions relate to chain_pairs destination rvs
+    chain_dest_rv_idxs: Set[int] = {pair[1].idx for pair in chain_pairs}
+    if not all(rv_idx in chain_dest_rv_idxs for rv_idx in initial_chain_conditioned_rvs.keys()):
+        raise ValueError('a chain initial condition is not a chain destination rv')
+
+    # Convert chain_pairs for registering with `sample_rvs`.
+    # rv_idx maps RandomVariable id to a position it exists in rvs (doesn't matter if rv is duplicated in rvs)
+    # copy_idx RandomVariable id to a position in rvs that it can be copied from for Markov chaining.
+    rv_idx: Dict[int, int] = {id(rv): i for i, rv in enumerate(rvs)}
+    copy_idx: Dict[int, int] = {id(rv): rv_idx[id(prev_rv)] for prev_rv, rv in chain_pairs}
+
+    # Get rv state slots, rvs_slots is co-indexed with rvs
+    slot_map: SlotMap = program_with_slotmap.slot_map
+    rvs_slots = tuple(tuple(slot_map[ind] for ind in rv) for rv in rvs)
+
+    sample_rvs: Sequence[SampleRV] = tuple(
+        SampleRV(idx, rv, rv_slots, copy_idx.get(id(rv)))
+        for idx, rv, rv_slots in zip(count(), rvs, rvs_slots)
+    )
+
+    # Process the condition to get zero and one slots
+    slots_0: Set[int] = set()
+    slots_1: Set[int] = set()
+    for rv in program_with_slotmap.rvs:
+        conditioning: Optional[Set[Indicator]] = conditioned_rvs.get(rv.idx)
+        if conditioning is not None:
+            slots_1.update(slot_map[ind] for ind in conditioning)
+            slots_0.update(slot_map[ind] for ind in rv if ind not in conditioning)
+            continue
+
+        conditioning: Optional[Set[Indicator]] = initial_chain_conditioned_rvs.get(rv.idx)
+        if conditioning is not None:
+            slots_1.update(slot_map[ind] for ind in conditioning)
+            slots_0.update(slot_map[ind] for ind in rv if ind not in conditioning)
+            continue
+
+        # default
+        slots_1.update(slot_map[ind] for ind in rv)
+
+    return SamplerInfo(
+        sample_rvs=sample_rvs,
+        condition=tuple(ind for condition_set in conditioned_rvs.values() for ind in condition_set),
+        yield_f=yield_f,
+        slots_0=slots_0,
+        slots_1=slots_1,
+    )
+
+
+def uniform_random_sample(
+        sample_rvs: Sequence[SampleRV],
+        slots_0: Collection[int],
+        slots_1: Collection[int],
+        slots: NDArrayNumeric,
+        state: NDArrayStates,
+        rand: Random,
+):
+    """
+    Helper for samplers.
+
+    Sets the states to a random instance and configures slots to match.
+    States are drawn from a uniform distribution, drawn using random.randrange.
+    """
+
+    # Set up the input slots to respect conditioning
+    for slot in slots_0:
+        slots[slot] = 0
+    for slot in slots_1:
+        slots[slot] = 1
+
+    for sample_rv in sample_rvs:
+        candidates = []
+        for slot_state, slot in enumerate(sample_rv.slots):
+            if slots[slot] == 1:
+                slots[slot] = 0
+                candidates.append((slot_state, slot))
+
+        # Pick a random state for sample_rv
+        slot_state, slot = candidates[rand.randrange(0, len(candidates))]
+        state[sample_rv.index] = slot_state
+        slots[slot] = 1

ck/sampling/uniform_sampler.py

@@ -0,0 +1,72 @@
+import random
+from typing import Set, List, Iterator, Optional, Sequence
+
+import numpy as np
+
+from ck.pgm import Instance, RandomVariable, Indicator
+from ck.probability.probability_space import dtype_for_state_indexes, Condition, check_condition
+from ck.utils.map_set import MapSet
+from ck.utils.np_extras import DType
+from ck.utils.random_extras import Random
+from .sampler import Sampler
+from .sampler_support import YieldF
+
+
+class UniformSampler(Sampler):
+
+    def __init__(
+            self,
+            rvs: RandomVariable | Sequence[RandomVariable],
+            condition: Condition = (),
+            rand: Random = random,
+    ):
+        condition: Sequence[Indicator] = check_condition(condition)
+
+        self._yield_f: YieldF
+        if isinstance(rvs, RandomVariable):
+            # a single rv
+            rvs = (rvs,)
+            self._yield_f = lambda x: x.item()
+        else:
+            # a sequence of rvs
+            self._yield_f = lambda x: x.tolist()
+
+        super().__init__(rvs, condition)
+
+        # Group condition indicators by `rv_idx`.
+        conditioned_rvs: MapSet[int, int] = MapSet()
+        for ind in condition:
+            conditioned_rvs.add(ind.rv_idx, ind.state_idx)
+
+        def get_possible_states(_rv: RandomVariable) -> List[int]:
+            """
+            Get the allowable states for a given random variable, given
+            conditions in `conditioned_rvs`.
+            """
+            condition_states: Optional[Set[int]] = conditioned_rvs.get(_rv.idx)
+            if condition_states is None:
+                return list(range(len(_rv)))
+            else:
+                return list(condition_states)
+
+        possible_states: List[List[int]] = [
+            get_possible_states(rv)
+            for rv in self.rvs
+        ]
+
+        self._possible_states: List[List[int]] = possible_states
+        self._rand: Random = rand
+        self._state_dtype: DType = dtype_for_state_indexes(self.rvs)
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        possible_states = self._possible_states
+        yield_f = self._yield_f
+        rand = self._rand
+        state = np.zeros(len(possible_states), dtype=self._state_dtype)
+        while True:
+            for i, l in enumerate(possible_states):
+                state_idx = rand.randrange(0, len(l))
+                state[i] = l[state_idx]
+            # We know the yield function will always provide either ints or Instances
+            # noinspection PyTypeChecker
+            yield yield_f(state)

ck/sampling/wmc_direct_sampler.py

@@ -0,0 +1,171 @@
+from typing import Collection, Iterator, Sequence
+
+import numpy as np
+
+from ck.pgm import Instance
+from ck.probability.probability_space import dtype_for_state_indexes
+from ck.program.program_buffer import ProgramBuffer
+from ck.program.raw_program import RawProgram
+from ck.sampling.sampler import Sampler
+from ck.sampling.sampler_support import SampleRV, YieldF, SamplerInfo
+from ck.utils.np_extras import NDArrayNumeric, NDArrayStates
+from ck.utils.random_extras import Random
+
+
+class WMCDirectSampler(Sampler):
+
+    def __init__(
+            self,
+            sampler_info: SamplerInfo,
+            raw_program: RawProgram,
+            rand: Random,
+    ):
+        super().__init__(sampler_info.rvs, sampler_info.condition)
+        self._yield_f: YieldF = sampler_info.yield_f
+        self._rand: Random = rand
+        self._program_buffer = ProgramBuffer(raw_program)
+        self._sample_rvs: Sequence[SampleRV] = tuple(sampler_info.sample_rvs)
+        self._chain_rvs: Sequence[SampleRV] = tuple(
+            sample_rv for sample_rv in sampler_info.sample_rvs if sample_rv.copy_index is not None)
+        self._state_dtype = dtype_for_state_indexes(self.rvs)
+        self._max_number_of_states: int = max((len(rv) for rv in self.rvs), default=0)
+        self._slots_1: Collection[int] = sampler_info.slots_1
+
+        # Set up the input slots to 0 or 1 to respect conditioning and initial Markov chain states.
+        slots: NDArrayNumeric = self._program_buffer.vars
+        for slot in sampler_info.slots_0:
+            slots[slot] = 0
+        for slot in sampler_info.slots_1:
+            slots[slot] = 1
+
+    def __iter__(self) -> Iterator[Instance] | Iterator[int]:
+        yield_f = self._yield_f
+        rand = self._rand
+        sample_rvs = self._sample_rvs
+        chain_rvs = self._chain_rvs
+        slots_1 = self._slots_1
+        program_buffer = self._program_buffer
+        slots: NDArrayNumeric = program_buffer.vars
+
+        # Calling wmc() will give the weighted model count for the state of the current input slots.
+        def wmc() -> float:
+            return program_buffer.compute().item()
+
+        # Set up working memory buffers
+        states: NDArrayStates = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+        buff_slots = np.zeros(self._max_number_of_states, dtype=np.uintp)
+        buff_states = np.zeros(self._max_number_of_states, dtype=self._state_dtype)
+
+        while True:
+            # Consider all possible instantiations given the conditions, c, where the instantiations are ordered.
+            # Let awmc(i|c) be the accumulated WMC of the ith instantiation.
+            # We want to find the smallest instantiation i such that
+            #     rnd <= awmc(i|c)
+            # where rnd is in [0, 1) * wmc().
+
+            rnd: float = rand.random() * wmc()
+
+            for sample_rv in sample_rvs:
+                # Prepare to loop over random variable states.
+                # Keep track of the non-zero slots in buff_slots and buff_states.
+                num_possible_states: int = 0
+                for j, slot in enumerate(sample_rv.slots):
+                    if slots[slot] != 0:
+                        buff_slots[num_possible_states] = slot
+                        buff_states[num_possible_states] = j
+                        num_possible_states += 1
+
+                if num_possible_states == 0:
+                    raise RuntimeError('zero probability')
+
+                # Try each possible state of the current random variable.
+                # Once a state is selected, then the following is true:
+                #   states[rv_position] = state
+                #   m_prev_states[rv_position] = state
+                #   slots set up to include condition rv = state.
+                #   rnd is reduced to account for the states skipped.
+                #
+                # We can do this either by sequentially checking each state or by doing
+                # a binary search. Here we start with binary search then finish sequentially
+                # once the candidates size falls below 'THRESHOLD'.
+
+                # Binary search
+                THRESHOLD = 2
+                lo: int = 0
+                hi: int = num_possible_states
+                w_0_mark: int = 0
+                w: float = 0
+                while lo + THRESHOLD < hi:
+                    mid: int = (lo + hi) // 2
+
+                    for i in range(mid, hi):
+                        slots[buff_slots[i]] = 0
+
+                    w = wmc()
+                    w_0_mark = mid
+                    if w < rnd:
+                        # wmc() is too low, the desired state is >= buff_states[mid]
+                        for i in range(mid, hi):
+                            slots[buff_slots[i]] = 1
+                        lo = mid
+                    else:
+                        # wmc() is too high, the desired state is < buff_states[mid]
+                        hi = mid
+
+                # Now the state we want is between lo (inclusive) and hi (exclusive).
+                # Slots at least up to lo will be set to 1.
+
+                # clear top slots, lo and up.
+                for k in range(lo, num_possible_states):
+                    slots[buff_slots[k]] = 0
+
+                # Adjust rnd to account for lo > 0.
+                if lo == 0:
+                    # The chances of this case may be low, but if so, then
+                    # slots[m_buff_slots[lo]] = 0  which implies wmc() == 0,
+                    # so we can save a call to wmc().
+                    pass
+                elif w_0_mark == lo:
+                    # We can use the last wmc() call, stored in w.
+                    # This saves a call to wmc().
+                    rnd -= w
+                else:
+                    rnd -= wmc()
+
+                # Clear remaining slots
+                for k in range(0, lo):
+                    slots[buff_slots[k]] = 0
+
+                # Sequential search
+                k = lo
+                while k < hi:
+                    slot = buff_slots[k]
+                    slots[slot] = 1
+                    w = wmc()
+                    if rnd < w:
+                        break
+                    slots[slot] = 0
+                    rnd -= w
+                    k += 1
+
+                slot = buff_slots[k]
+                state = buff_states[k]
+                slots[slot] = 1
+                states[sample_rv.index] = state
+
+            # We know the yield function will always provide either ints or Instances
+            # noinspection PyTypeChecker
+            yield yield_f(states)
+
+            # Reset the one slots for the next iteration.
+            for slot in slots_1:
+                slots[slot] = 1
+
+            # Copy chain pairs for next iteration.
+            # (This writes over any initial chain conditions from slots_1.)
+            for sample_rv in chain_rvs:
+                rv_slots = sample_rv.slots
+                prev_state_idx: int = states.item(sample_rv.copy_index)
+                for slot in rv_slots:
+                    slots[slot] = 0
+                slots[rv_slots[prev_state_idx]] = 1