compiled-knowledge 4.0.0a9__cp312-cp312-win_amd64.whl → 4.0.0a11__cp312-cp312-win_amd64.whl

ck/pgm_circuit/pgm_circuit.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Sequence, List, Dict
 
-from ck.circuit import CircuitNode
+from ck.circuit import CircuitNode, Circuit
 from ck.pgm import RandomVariable, Indicator
 from ck.pgm_circuit.slot_map import SlotMap, SlotKey
 from ck.utils.np_extras import NDArray
@@ -46,19 +46,21 @@ class PGMCircuit:
 
     def dump(self, *, prefix: str = '', indent: str = '    ') -> None:
         """
-        Print a dump of the Join Tree.
+        Print a dump of the circuit.
         This is intended for debugging and demonstration purposes.
 
-        Each cluster is printed as: {separator rvs} | {non-separator rvs}.
-
         Args:
             prefix: optional prefix for indenting all lines.
             indent: additional prefix to use for extra indentation.
         """
-        # Name the circuit variables
-        circuit = self.circuit_top.circuit
+
+        # We infer names for the circuit variables, either as an indicator or as a parameter.
+        # The `var_names` will be passed to `circuit.dump`.
+
+        circuit: Circuit = self.circuit_top.circuit
         var_names: List[str] = [''] * circuit.number_of_vars
 
+        # Name the circuit variables that are indicators
         rvs_by_idx: Dict[int, RandomVariable] = {rv.idx: rv for rv in self.rvs}
         slot_key: SlotKey
         slot: int
@@ -68,8 +70,10 @@ class PGMCircuit:
                 state_idx = slot_key.state_idx
                 var_names[slot] = f'{rv.name!r}[{state_idx}] {rv.states[state_idx]!r}'
 
-        for slot, param_value in enumerate(self.parameter_values, start=self.number_of_indicators):
-            var_names[slot] = f'param {param_value}'
+        # Name the circuit variables that are parameters
+        for i, param_value in enumerate(self.parameter_values):
+            slot = i + self.number_of_indicators
+            var_names[slot] = f'param[{i}] = {param_value}'
 
-        # Show all the slots
+        # Dump the circuit
         circuit.dump(prefix=prefix, indent=indent, var_names=var_names)

ck/pgm_circuit/program_with_slotmap.py

@@ -1,4 +1,4 @@
-from typing import Tuple, Iterator, Sequence, Dict, Iterable
+from typing import Tuple, Sequence, Dict, Iterable
 
 from ck.pgm import RandomVariable, rv_instances, Instance, rv_instances_as_indicators, Indicator, ParamId
 from ck.pgm_circuit.slot_map import SlotMap, SlotKey
@@ -30,11 +30,13 @@ class ProgramWithSlotmap:
                 has a length and rv[i] is a unique 'indicator' across all rvs.
             precondition: conditions on rvs that are compiled into the program.
 
+        Raises:
+            ValueError: if rvs contains duplicates.
         """
         self._program_buffer: ProgramBuffer = program_buffer
         self._slot_map: SlotMap = slot_map
         self._rvs: Tuple[RandomVariable, ...] = tuple(rvs)
-        self._precondition: Sequence[Indicator] = precondition
+        self._precondition: Tuple[Indicator, ...] = tuple(precondition)
 
         if len(rvs) != len(set(rv.idx for rv in rvs)):
             raise ValueError('duplicate random variables provided')
@@ -67,7 +69,7 @@ class ProgramWithSlotmap:
     def slot_map(self) -> SlotMap:
         return self._slot_map
 
-    def instances(self, flip: bool = False) -> Iterator[Instance]:
+    def instances(self, flip: bool = False) -> Iterable[Instance]:
         """
         Enumerate instances of the random variables.
 
@@ -84,7 +86,7 @@ class ProgramWithSlotmap:
         """
         return rv_instances(*self._rvs, flip=flip)
 
-    def instances_as_indicators(self, flip: bool = False) -> Iterator[Sequence[Indicator]]:
+    def instances_as_indicators(self, flip: bool = False) -> Iterable[Sequence[Indicator]]:
         """
         Enumerate instances of the random variables.
 

ck/pgm_compiler/ace/ace.py

@@ -5,7 +5,9 @@ from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional, List, Tuple
 
-from ck.circuit import CircuitNode
+import numpy as np
+
+from ck.circuit import CircuitNode, Circuit
 from ck.in_out.parse_ace_lmap import read_lmap, LiteralMap
 from ck.in_out.parse_ace_nnf import read_nnf_with_literal_map
 from ck.in_out.render_net import render_bayesian_network
@@ -58,6 +60,22 @@ def compile_pgm(
     if check_is_bayesian_network and not pgm.check_is_bayesian_network():
         raise ValueError('the given PGM is not a Bayesian network')
 
+    # ACE cannot deal with the empty PGM even though it is a valid Bayesian network
+    if pgm.number_of_factors == 0:
+        circuit = Circuit()
+        circuit.new_vars(pgm.number_of_indicators)
+        parameter_values = np.array([], dtype=np.float64)
+        slot_map = {indicator: i for i, indicator in enumerate(pgm.indicators)}
+        return PGMCircuit(
+            rvs=pgm.rvs,
+            conditions=(),
+            circuit_top=circuit.const(1),
+            number_of_indicators=pgm.number_of_indicators,
+            number_of_parameters=0,
+            slot_map=slot_map,
+            parameter_values=parameter_values,
+        )
+
     java: str
     classpath_separator: str
     java, classpath_separator = _find_java()
@@ -69,7 +87,7 @@ def compile_pgm(
     )
     ace_cmd: List[str] = [
         java,
-        f'-cp',
+        '-cp',
         class_path,
         f'-DACEC2D={files.c2d}',
         f'-Xmx{int(m_bytes)}m',
@@ -83,7 +101,7 @@ def compile_pgm(
             node_names: List[str] = render_bayesian_network(pgm, file, check_structure_bayesian=False)
 
         # Run Ace
-        ace_result = subprocess.run(ace_cmd, capture_output=True, text=True)
+        ace_result: subprocess.CompletedProcess = subprocess.run(ace_cmd, capture_output=True, text=True)
         if print_output:
             print(ace_result.stdout)
             print(ace_result.stderr)
@@ -126,6 +144,29 @@ def compile_pgm(
     )
 
 
+def ace_available(
+        ace_dir: Optional[Path | str] = None,
+        jar_dir: Optional[Path | str] = None,
+) -> bool:
+    """
+    Returns:
+        True if it looks like ACE is available, False otherwise.
+        ACE is available if ACE files are in the default location and Java is available.
+    """
+    try:
+        java: str
+        java, _ = _find_java()
+        _: _AceFiles = _find_ace_files(ace_dir, jar_dir)
+
+        java_cmd: List[str] = [java, '--version',]
+        java_result: subprocess.CompletedProcess = subprocess.run(java_cmd, capture_output=True, text=True)
+
+        return java_result.returncode == 0
+
+    except RuntimeError:
+        return False
+
+
 def copy_ace_to_default_location(
         ace_dir: Path | str,
         jar_dir: Optional[Path | str] = None,
@@ -179,6 +220,9 @@ def _find_java() -> Tuple[str, str]:
 
     Returns:
         (java, classpath_separator)
+
+    Raises:
+        RuntimeError: if not found, including a helpful message.
     """
     if sys.platform == 'win32':
         return 'java.exe', ';'
@@ -198,7 +242,7 @@ def _find_ace_files(
     Look for the needed Ace files.
 
     Raises:
-        RuntimeError: if not found, including a helpful message .
+        RuntimeError: if not found, including a helpful message.
     """
     ace_dir: Path = default_ace_location() if ace_dir is None else Path(ace_dir)
     jar_dir: Path = ace_dir if jar_dir is None else Path(jar_dir)

ck/pgm_compiler/factor_elimination.py

@@ -12,19 +12,21 @@ from ck.pgm_compiler.support.join_tree import *
 
 _NEG_INF = float('-inf')
 
+DEFAULT_PRODUCT_SEARCH_LIMIT: int = 1000
+
 
 def compile_pgm(
         pgm: PGM,
         const_parameters: bool = True,
         *,
         algorithm: JoinTreeAlgorithm = MIN_FILL_THEN_DEGREE,
-        limit_product_tree_search: int = 1000,
+        limit_product_tree_search: int = DEFAULT_PRODUCT_SEARCH_LIMIT,
         pre_prune_factor_tables: bool = True,
 ) -> PGMCircuit:
     """
     Compile the PGM to an arithmetic circuit, using factor elimination.
 
-    When forming the product of factors withing a join tree nodes,
+    When forming the product of factors within join tree nodes,
     this method searches all practical binary trees for forming products,
     up to the given limit, `limit_product_tree_search`. The minimum is 1.
 
@@ -57,7 +59,7 @@ def compile_pgm_best_jointree(
         pgm: PGM,
         const_parameters: bool = True,
         *,
-        limit_product_tree_search: int = 1000,
+        limit_product_tree_search: int = DEFAULT_PRODUCT_SEARCH_LIMIT,
         pre_prune_factor_tables: bool = True,
 ) -> PGMCircuit:
     """
@@ -111,7 +113,7 @@ def compile_pgm_best_jointree(
 def join_tree_to_circuit(
         join_tree: JoinTree,
         const_parameters: bool = True,
-        limit_product_tree_search: int = 1000,
+        limit_product_tree_search: int = DEFAULT_PRODUCT_SEARCH_LIMIT,
         pre_prune_factor_tables: bool = True,
 ) -> PGMCircuit:
     """

ck/pgm_compiler/recursive_conditioning.py

@@ -52,10 +52,15 @@ def compile_pgm(
         multiply_indicators=True,
         pre_prune_factor_tables=pre_prune_factor_tables,
     )
-    dtree: _DTree = _make_dtree(elimination_order, factor_tables)
 
-    states: List[Sequence[int]] = [tuple(range(len(rv))) for rv in pgm.rvs]
-    top: CircuitNode = dtree.make_circuit(states, factor_tables.circuit)
+    if pgm.number_of_factors == 0:
+        # Deal with special case: no factors
+        top: CircuitNode = factor_tables.circuit.const(1)
+    else:
+        dtree: _DTree = _make_dtree(elimination_order, factor_tables)
+        states: List[Sequence[int]] = [tuple(range(len(rv))) for rv in pgm.rvs]
+        top: CircuitNode = dtree.make_circuit(states, factor_tables.circuit)
+
     top.circuit.remove_unreachable_op_nodes(top)
 
     return PGMCircuit(

ck/pgm_compiler/support/clusters.py

@@ -9,7 +9,7 @@ from ck.pgm import PGM
 
 # A VEObjective is a variable elimination objective function.
 # An objective function is a function from a random variable index (int)
-# to an objected value (float or int). This is used to select
+# to an objective value (float or int). This is used to select
 # a random variable to eliminate in `ve_greedy_min`.
 VEObjective = Callable[[int], int | float]
 

ck/pgm_compiler/variable_elimination.py

@@ -69,13 +69,13 @@ def compile_pgm(
                 tables_with_rv.append(product(x, y))
             next_tables.append(sum_out(tables_with_rv[0], (rv_idx,)))
             cur_tables = next_tables
-    # All rvs are now eliminated
+
+    # All rvs are now eliminated - all tables should have a single top.
     tops: List[CircuitNode] = [
         table.top()
         for table in cur_tables
-        if len(table) > 0
     ]
-    top = factor_tables.circuit.optimised_add(tops)
+    top: CircuitNode = factor_tables.circuit.optimised_mul(tops)
     top.circuit.remove_unreachable_op_nodes(top)
 
     return PGMCircuit(

ck/probability/empirical_probability_space.py

@@ -7,6 +7,9 @@ from ck.probability.probability_space import ProbabilitySpace, Condition, check_
 class EmpiricalProbabilitySpace(ProbabilitySpace):
     def __init__(self, rvs: Sequence[RandomVariable], samples: Iterable[Instance]):
         """
+        Enable probabilistic queries over a sample from a sample space.
+        Note that this is not necessarily an efficient approach to calculating probabilities and statistics.
+
         Assumes:
             len(sample) == len(rvs), for each sample in samples.
             0 <= sample[i] < len(rvs[i]), for each sample in samples, for i in range(len(rvs)).

ck/probability/pgm_probability_space.py

@@ -0,0 +1,32 @@
+from typing import Sequence, Iterable, Tuple, Dict, List
+
+from ck.pgm import RandomVariable, Indicator, Instance, PGM
+from ck.probability.probability_space import ProbabilitySpace, Condition, check_condition
+
+
+class PGMProbabilitySpace(ProbabilitySpace):
+    def __init__(self, pgm: PGM):
+        """
+        Enable probabilistic queries directly on a PGM.
+        Note that this is not necessarily an efficient approach to calculating probabilities and statistics.
+
+        Args:
+            pgm: The PGM to query.
+        """
+        self._pgm = pgm
+        self._z = None
+
+    @property
+    def rvs(self) -> Sequence[RandomVariable]:
+        return self._pgm.rvs
+
+    def wmc(self, *condition: Condition) -> float:
+        condition: Tuple[Indicator, ...] = check_condition(condition)
+        return self._pgm.value_product_indicators(*condition)
+
+    @property
+    def z(self) -> float:
+        if self._z is None:
+            self._z = self._pgm.value_product_indicators()
+        return self._z
+

ck/probability/probability_space.py

@@ -3,6 +3,7 @@ An abstract class for object providing probabilities.
 """
 import math
 from abc import ABC, abstractmethod
+from itertools import chain
 from typing import Sequence, Tuple, Iterable, Callable
 
 import numpy as np
@@ -75,15 +76,38 @@ class ProbabilitySpace(ABC):
             the probability of the given indicators, conditioned on the given conditions.
         """
         condition: Tuple[Indicator, ...] = check_condition(condition)
+
         if len(condition) == 0:
             z = self.z
+            if z <= 0:
+                return np.nan
         else:
             z = self.wmc(*condition)
-            indicators += condition  # append the condition
-        if z > 0:
-            return self.wmc(*indicators) / z
-        else:
-            return np.nan
+            if z <= 0:
+                return np.nan
+
+            # Combine the indicators with the condition
+            # If a variable is mentioned in both the indicators and condition, then
+            # we need to take the intersection, and check for contradictions.
+            # If a variable is mentioned in the condition but not indicators, then
+            # the rv condition needs to be added to the indicators.
+            indicator_groups: MapSet[int, Indicator] = _group_indicators(indicators)
+            condition_groups: MapSet[int, Indicator] = _group_indicators(condition)
+
+            for rv_idx, indicators in condition_groups.items():
+                indicator_group = indicator_groups.get(rv_idx)
+                if indicator_group is None:
+                    indicator_groups.add_all(rv_idx, indicators)
+                else:
+                    indicator_group.intersection_update(indicators)
+                    if len(indicator_group) == 0:
+                        # A contradiction between the indicators and conditions
+                        return 0.0
+
+            # Collect all the indicators from the updated indicator_groups
+            indicators = chain(*indicator_groups.values())
+
+        return self.wmc(*indicators) / z
 
     def marginal_distribution(self, *rvs: RandomVariable, condition: Condition = ()) -> NDArrayNumeric:
         """
@@ -160,9 +184,7 @@ class ProbabilitySpace(ABC):
         assert len(rv_indexes) == len(rvs), 'duplicated random variables not allowed'
 
         # Group conditioning indicators by random variable.
-        conditions_by_rvs = MapSet()
-        for ind in condition:
-            conditions_by_rvs.get_set(ind.rv_idx).add(ind.state_idx)
+        conditions_by_rvs = _group_states(condition)
 
         # See if any MAP random variable is also conditioned.
         # Reduce the state space of any conditioned MAP rv.
@@ -195,12 +217,12 @@ class ProbabilitySpace(ABC):
             # Loop over the state space of the 'loop' rvs
             best_probability = float('-inf')
             best_states = None
-            inds: Tuple[Indicator, ...]
-            for inds in _combos(loop_rvs):
-                probability = self.wmc(*(inds + new_conditions))
+            indicators: Tuple[Indicator, ...]
+            for indicators in _combos(loop_rvs):
+                probability = self.wmc(*(indicators + new_conditions))
                 if probability > best_probability:
                     best_probability = probability
-                    best_states = tuple(ind.state_idx for ind in inds)
+                    best_states = tuple(ind.state_idx for ind in indicators)
             condition_probability = self.wmc(*condition)
             return best_probability / condition_probability, best_states
 
@@ -551,6 +573,38 @@ def dtype_for_state_indexes(rvs: Iterable[RandomVariable]) -> DTypeStates:
     return dtype_for_number_of_states(max((len(rv) for rv in rvs), default=0))
 
 
+def _group_indicators(indicators: Iterable[Indicator]) -> MapSet[int, Indicator]:
+    """
+    Group the given indicators by rv_idx.
+
+    Args:
+        indicators: the indicators to group.
+
+    Returns:
+        A mapping from rv_idx to set of indicators.
+    """
+    groups: MapSet[int, Indicator] = MapSet()
+    for indicator in indicators:
+        groups.add(indicator.rv_idx, indicator)
+    return groups
+
+
+def _group_states(indicators: Iterable[Indicator]) -> MapSet[int, int]:
+    """
+    Group the given indicator states by rv_idx.
+
+    Args:
+        indicators: the indicators to group.
+
+    Returns:
+        A mapping from rv_idx to set of state indexes.
+    """
+    groups: MapSet[int, int] = MapSet()
+    for indicator in indicators:
+        groups.add(indicator.rv_idx, indicator.state_idx)
+    return groups
+
+
 def _normalise_marginal(distribution: NDArrayFloat64) -> None:
     """
     Update the values in the given distribution to

ck/program/program.py

@@ -1,3 +1,6 @@
+"""
+For more documentation on this module, refer to the Jupyter notebook docs/6_circuits_and_programs.ipynb.
+"""
 from typing import Callable, Sequence
 
 import numpy as np
@@ -8,7 +11,12 @@ from ck.utils.np_extras import DTypeNumeric, NDArrayNumeric
 
 class Program:
     """
-    A Program wraps a RawProgram to make a convenient callable.
+    A program represents an arithmetic a function from input values to output values.
+
+    Internally a `Program` wraps a `RawProgram` which is the object returned by a circuit compiler.
+
+    Every `Program` has a numpy `dtype` which defines the numeric data type for input and output values.
+    Typically, the `dtype` of a program is a C style double.
     """
 
     def __init__(self, raw_program: RawProgram):

ck/program/raw_program.py

@@ -9,8 +9,8 @@ from ck.utils.np_extras import NDArrayNumeric, DTypeNumeric
 
 # RawProgramFunction is a function of three ctypes arrays, returning nothing.
 # Args:
-#     [0]: input parameter values,
-#     [1]: working memory,
+#     [0]: input  values,
+#     [1]: temporary working memory,
 #     [2]: output values.
 RawProgramFunction = Callable[[ct.POINTER, ct.POINTER, ct.POINTER], None]
 
@@ -18,10 +18,16 @@ RawProgramFunction = Callable[[ct.POINTER, ct.POINTER, ct.POINTER], None]
 @dataclass
 class RawProgram:
     """
+    A raw program is returned by a circuit compiler to provide execution of
+    the function defined by a compiled circuit.
+
+    A `RawProgram` is a `Callable` with the signature:
+
+
     Fields:
         function: is a function of three ctypes arrays, returning nothing.
         dtype: the numpy data type of  the array values.
-        number_of_vars: the number of input parameter values (first function argument).
+        number_of_vars: the number of input values (first function argument).
         number_of_tmps: the number of working memory values (second function argument).
         number_of_results: the number of result values (third function argument).
         var_indices:  maps the index of inputs (from 0 to self.number_of_vars - 1) to the index

ck/sampling/sampler_support.py

@@ -13,7 +13,7 @@ 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], Instance | int]
+YieldF = Callable[[NDArrayStates], int] | Callable[[NDArrayStates], Instance]
 
 
 @dataclass

ck/sampling/uniform_sampler.py

@@ -1,15 +1,15 @@
-from typing import Set, List, Iterator, Optional, Sequence
 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 .sampler import Sampler
-from .sampler_support import YieldF
 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):
@@ -39,11 +39,15 @@ class UniformSampler(Sampler):
             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 [state_idx for state_idx in range(len(_rv)) if state_idx not in condition_states]
+                return list(condition_states)
 
         possible_states: List[List[int]] = [
             get_possible_states(rv)
@@ -63,4 +67,6 @@ class UniformSampler(Sampler):
             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

@@ -8,7 +8,7 @@ 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
+from ck.utils.np_extras import NDArrayNumeric, NDArrayStates
 from ck.utils.random_extras import Random
 
 
@@ -52,7 +52,7 @@ class WMCDirectSampler(Sampler):
             return program_buffer.compute().item()
 
         # Set up working memory buffers
-        states = np.zeros(len(sample_rvs), dtype=self._state_dtype)
+        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)
 
@@ -153,6 +153,8 @@ class WMCDirectSampler(Sampler):
                 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.

ck/sampling/wmc_gibbs_sampler.py

@@ -65,12 +65,16 @@ class WMCGibbsSampler(Sampler):
             if skip == 0:
                 while True:
                     self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
                     yield yield_f(state)
             else:
                 while True:
                     for _ in range(skip):
                         self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
                     self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
                     yield yield_f(state)
 
         else:
@@ -79,6 +83,8 @@ class WMCGibbsSampler(Sampler):
                 for _ in range(skip):
                     self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
                 self._next_sample_gibbs(sample_rvs, slots_1, program_buffer, prs, state, rand)
+                # We know the yield function will always provide either ints or Instances
+                # noinspection PyTypeChecker
                 yield yield_f(state)
                 if rand.random() < pr_restart:
                     # Set an initial system state

ck/sampling/wmc_metropolis_sampler.py

@@ -1,4 +1,4 @@
-from typing import Collection, Iterator, List, Sequence
+from typing import Collection, Iterator, Sequence
 
 import numpy as np
 
@@ -78,12 +78,16 @@ class WMCMetropolisSampler(Sampler):
             if skip == 0:
                 while True:
                     w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
                     yield yield_f(state)
             else:
                 while True:
                     for _ in range(skip):
                         w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
                     w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                    # We know the yield function will always provide either ints or Instances
+                    # noinspection PyTypeChecker
                     yield yield_f(state)
 
         else:
@@ -92,6 +96,8 @@ class WMCMetropolisSampler(Sampler):
                 for _ in range(skip):
                     w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
                 w = self._next_sample_metropolis(possibles, program_buffer, state, w, rand)
+                # We know the yield function will always provide either ints or Instances
+                # noinspection PyTypeChecker
                 yield yield_f(state)
 
                 if rand.random() < pr_restart:

ck/sampling/wmc_rejection_sampler.py

@@ -91,6 +91,8 @@ class WMCRejectionSampler(Sampler):
             w: float = wmc()
 
             if rand.random() * self._w_max < w:
+                # We know the yield function will always provide either ints or Instances
+                # noinspection PyTypeChecker
                 yield yield_f(state)
 
             # Update w_not_seen and w_high to adapt w_max.