Models

All models are torch.nn.Modules with the same learning interface:

• model.update(x, y) -> votes — one learning step on a batch,
• model(x) / model.forward(x) -> votes — vote sums (B, n_outputs),
• model.predict(x), model.predict_proba(x), model.confidence(x),
• model.rules(), model.clause_expression(j), model.included_mask(), model.state_summary().

Dense (segmentation) models add model.vote_map(x) -> (B, K, H, W), model.predict(x) -> (B, H, W) and model.clause_map(x); their forward() returns the folded votes (B·H·W, K).

GraphTsetlinMachine takes a torchtsetlin.data.GraphBatch instead of a tensor and adds model.node_clause_outputs(g) and model.node_messages(g, layer) (see the graph guide).

torchtsetlin.models.TsetlinMachine

TsetlinMachine(n_features: Optional[int], n_classes: int, n_clauses: int, T: float, s: float = 10.0, *, weighted: bool = False, max_weight: Optional[int] = None, focused_negative_sampling: bool = False, **kwargs)

Bases: TsetlinMachineBase

Multi-class Tsetlin machine with class-owned clauses of positive/negative polarity.

Each of the n_classes classes owns n_clauses clauses. The first half of them vote for the class (positive polarity), the second half vote against it (negative polarity). The class sum is v_k = sum(w_j c_j) - sum(w_j c_j) over the two halves, clamped to [-T, T], and the prediction is argmax_k v_k.

Learning follows the standard algorithm (Granmo, 2018): for an example of class y the clauses of y receive feedback with probability (T - v_y)/(2T) (Type I to positive clauses, Type II to negative ones) and the clauses of one randomly chosen other class q receive feedback with probability (T + v_q)/(2T) (Type II to positive, Type I to negative clauses).

Parameters:

Name	Type	Description	Default
n_features	Optional[int]	number of Boolean features (None = infer from the first batch).	required
n_classes	int	number of classes (>= 2). Labels are integers 0..n_classes-1.	required
n_clauses	int	clauses per class.	required
T	float	vote margin (threshold). Larger T makes more clauses cooperate per class.	required
s	float	specificity.	10.0
weighted	bool	learn an integer weight per clause (Integer-Weighted TM). Weights start at 1, grow with Type Ia feedback and shrink with Type II feedback (min 1, as in pyTsetlinMachine and tmu).	False
max_weight	Optional[int]	optional upper bound on the clause weights (PyTsetlinMachineCUDA stores them as uint8, so its weights stop at 255). None: unbounded.	None
focused_negative_sampling	bool	sample the negative class proportionally to how strongly it currently votes, instead of uniformly.	False
**kwargs	Any	forwarded to :class:~torchtsetlin.models.TsetlinMachineBase (n_states, boost_true_positive, max_included_literals, drop_clause_p, drop_literal_p, init, state_dtype, ...).	{}

Shape

• input x: (B, F) Boolean (extra trailing dimensions are flattened).
• output of :meth:forward: (B, n_classes) float vote sums in [-T, T].

Example

model = TsetlinMachine(n_features=12, n_classes=2, n_clauses=20, T=10, s=3.9) votes = model.update(x, y) # learn from a batch, returns pre-update votes pred = model(x).argmax(dim=1) # predict

Source code in src/torchtsetlin/models/classifier.py

def __init__(
    self,
    n_features: Optional[int],
    n_classes: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    weighted: bool = False,
    max_weight: Optional[int] = None,
    focused_negative_sampling: bool = False,
    **kwargs,
) -> None:
    if n_classes < 2:
        raise ValueError("n_classes must be >= 2 (use two classes for binary problems)")
    if n_clauses < 2:
        raise ValueError("n_clauses (per class) must be >= 2")
    if T <= 0:
        raise ValueError("T must be positive")
    self.n_classes = int(n_classes)
    self.n_clauses = int(n_clauses)
    self.T = float(T)
    self.weighted = bool(weighted)
    if max_weight is not None and int(max_weight) < 1:
        raise ValueError("max_weight must be >= 1")
    self.max_weight = None if max_weight is None else int(max_weight)
    self.focused_negative_sampling = bool(focused_negative_sampling)
    super().__init__(n_features, self.n_classes * self.n_clauses, s=s, **kwargs)

    C = self.n_clauses_total
    clause_class = torch.arange(C) // self.n_clauses
    within = torch.arange(C) % self.n_clauses
    n_pos = (self.n_clauses + 1) // 2
    polarity = torch.where(within < n_pos, 1.0, -1.0)
    self.register_buffer("clause_class", clause_class, persistent=False)
    self.register_buffer("clause_polarity", polarity, persistent=False)
    self.register_buffer("weights", torch.ones(C, dtype=torch.int32))
    self.class_names: Optional[List[str]] = None

predict

predict(x) -> Tensor

Predicted class index per example (B,).

Source code in src/torchtsetlin/models/classifier.py

def predict(self, x) -> Tensor:
    """Predicted class index per example ``(B,)``."""
    return self.forward(x).argmax(dim=1)

predict_proba

predict_proba(x, method: str = 'linear', temperature: float = 1.0) -> Tensor

Class probabilities (B, n_classes).

method="linear" uses the calibrated score of Abeyrathna et al., p_k ∝ (1 + v_k/T)/2; method="softmax" applies a softmax to v/T.

Source code in src/torchtsetlin/models/classifier.py

def predict_proba(self, x, method: str = "linear", temperature: float = 1.0) -> Tensor:
    """Class probabilities ``(B, n_classes)``.

    ``method="linear"`` uses the calibrated score of Abeyrathna et al.,
    ``p_k ∝ (1 + v_k/T)/2``; ``method="softmax"`` applies a softmax to ``v/T``.
    """
    votes = self.forward(x).clamp(-self.T, self.T)  # probabilities live on [-T, T]
    if method == "softmax":
        return F.predict_proba_from_votes(votes, self.T, temperature)
    p = 0.5 * (1.0 + votes / self.T)
    return p / p.sum(dim=1, keepdim=True).clamp_min(1e-12)

confidence

confidence(x) -> Tensor

Confidence of the winning class in [0, 1]: (v_max + T) / (2T).

Source code in src/torchtsetlin/models/classifier.py

def confidence(self, x) -> Tensor:
    """Confidence of the winning class in ``[0, 1]``: ``(v_max + T) / (2T)``."""
    return F.confidence_from_votes(self.forward(x), self.T)

clause_info

clause_info(clause: int, feature_names: Optional[Sequence[str]] = None) -> Dict

Metadata for one clause: class, polarity, weight and the literal expression.

Source code in src/torchtsetlin/models/classifier.py

def clause_info(self, clause: int, feature_names: Optional[Sequence[str]] = None) -> Dict:
    """Metadata for one clause: class, polarity, weight and the literal expression."""
    cls = int(self.clause_class[clause])
    pol = int(self.clause_polarity[clause])
    name = self.class_names[cls] if self.class_names else str(cls)
    return {
        "clause": int(clause),
        "class": cls,
        "class_name": name,
        "polarity": pol,
        "weight": int(self.weights[clause]) if self.weighted else 1,
        "n_literals": int(self.include_count[clause]),
        "expression": self.clause_expression(clause, feature_names),
    }

rules

rules(feature_names: Optional[Sequence[str]] = None, class_names: Optional[Sequence[str]] = None, skip_empty: bool = True, polarity: Optional[int] = None, class_index: Optional[int] = None) -> List[str]

Human readable IF ... THEN rules for all (non-empty) clauses.

Source code in src/torchtsetlin/models/classifier.py

def rules(
    self,
    feature_names: Optional[Sequence[str]] = None,
    class_names: Optional[Sequence[str]] = None,
    skip_empty: bool = True,
    polarity: Optional[int] = None,
    class_index: Optional[int] = None,
) -> List[str]:
    """Human readable ``IF ... THEN`` rules for all (non-empty) clauses."""
    names = list(class_names or self.class_names or [str(k) for k in range(self.n_classes)])
    out = []
    for j in range(self.n_clauses_total):
        if skip_empty and int(self.include_count[j]) == 0:
            continue
        pol = int(self.clause_polarity[j])
        if polarity is not None and pol != polarity:
            continue
        k = int(self.clause_class[j])
        if class_index is not None and k != class_index:
            continue
        expr = self.clause_expression(j, feature_names)
        w = f" (weight {int(self.weights[j])})" if self.weighted else ""
        verdict = f"THEN {names[k]}" if pol > 0 else f"THEN NOT {names[k]}"
        out.append(f"IF {expr} {verdict}{w}")
    return out

torchtsetlin.models.CoalescedTsetlinMachine

CoalescedTsetlinMachine(n_features: Optional[int], n_outputs: int, n_clauses: int, T: float, s: float = 10.0, *, multi_label: bool = False, negative_scale: float = 1.0, max_weight: Optional[int] = None, focused_negative_sampling: bool = False, **kwargs)

Bases: TsetlinMachineBase

Coalesced (multi-output) Tsetlin machine with clause sharing (Glimsdal & Granmo, 2021).

All n_clauses clauses are shared by the n_outputs outputs. Clause j votes for output i with an integer weight W[j, i]: positive weights vote for, negative weights against. Vote sums are v = c @ W clamped to [-T, T].

Learning per output i with binary target y_i: with probability d_i = |q_i - v_i| / (2T) (q_i = T if y_i = 1 else -T) a clause receives Type I feedback when its weight sign agrees with y_i and Type II feedback otherwise. Weights of matching, selected clauses move by +1 for y_i = 1 and -1 for y_i = 0, so a clause can change sides over time.

In multi-class mode (multi_label=False) the target class and one random negative class are updated per example; in multi-label mode every output is updated with its own probability (negative outputs scaled by negative_scale).

Parameters:

Name	Type	Description	Default
n_features	Optional[int]	Boolean feature count (None = lazy).	required
n_outputs	int	number of classes (multi-class) or labels (multi-label).	required
n_clauses	int	total number of shared clauses.	required
T	float	vote margin.	required
s	float	specificity.	10.0
multi_label	bool	interpret y as a (B, n_outputs) binary matrix and predict v_i > 0 per output.	False
max_weight	Optional[int]	optional bound on the magnitude of the signed clause weights (None: unbounded).	None
negative_scale	float	multiplier of the Type II probability for negative outputs in multi-label mode (the q hyper-parameter of the multi-task CTM literature).	1.0
focused_negative_sampling	bool	multi-class only; bias the negative class towards strongly voting classes.	False
**kwargs	Any	forwarded to :class:~torchtsetlin.models.TsetlinMachineBase.	{}

Source code in src/torchtsetlin/models/coalesced.py

def __init__(
    self,
    n_features: Optional[int],
    n_outputs: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    multi_label: bool = False,
    negative_scale: float = 1.0,
    max_weight: Optional[int] = None,
    focused_negative_sampling: bool = False,
    **kwargs,
) -> None:
    if n_outputs < 1 or (not multi_label and n_outputs < 2):
        raise ValueError("n_outputs must be >= 2 for multi-class, >= 1 for multi-label")
    if T <= 0:
        raise ValueError("T must be positive")
    self.n_outputs = int(n_outputs)
    self.n_clauses = int(n_clauses)
    self.T = float(T)
    self.multi_label = bool(multi_label)
    self.negative_scale = float(negative_scale)
    if max_weight is not None and int(max_weight) < 1:
        raise ValueError("max_weight must be >= 1")
    self.max_weight = None if max_weight is None else int(max_weight)
    self.focused_negative_sampling = bool(focused_negative_sampling)
    super().__init__(n_features, self.n_clauses, s=s, **kwargs)
    w = torch.randint(0, 2, (self.n_clauses, self.n_outputs), dtype=torch.int32) * 2 - 1
    self.register_buffer("weights", w)
    self.class_names: Optional[List[str]] = None

predict

predict(x) -> Tensor

Class indices (B,) (multi-class) or a Boolean (B, n_outputs) matrix.

Source code in src/torchtsetlin/models/coalesced.py

def predict(self, x) -> Tensor:
    """Class indices ``(B,)`` (multi-class) or a Boolean ``(B, n_outputs)`` matrix."""
    votes = self.forward(x)
    if self.multi_label:
        return votes > 0
    return votes.argmax(dim=1)

predict_proba

predict_proba(x, method: str = 'linear', temperature: float = 1.0) -> Tensor

Probabilities: per-output (1 + v/T)/2 (multi-label) or normalised over classes.

Source code in src/torchtsetlin/models/coalesced.py

def predict_proba(self, x, method: str = "linear", temperature: float = 1.0) -> Tensor:
    """Probabilities: per-output ``(1 + v/T)/2`` (multi-label) or normalised over classes."""
    votes = self.forward(x).clamp(-self.T, self.T)  # probabilities live on [-T, T]
    p = 0.5 * (1.0 + votes / self.T)
    if self.multi_label:
        return p
    if method == "softmax":
        return F.predict_proba_from_votes(votes, self.T, temperature)
    return p / p.sum(dim=1, keepdim=True).clamp_min(1e-12)

rules

rules(feature_names: Optional[Sequence[str]] = None, class_names: Optional[Sequence[str]] = None, skip_empty: bool = True, min_abs_weight: int = 1) -> List[str]

IF ... THEN rules listing every output the clause votes for/against with its weight.

Source code in src/torchtsetlin/models/coalesced.py

def rules(
    self,
    feature_names: Optional[Sequence[str]] = None,
    class_names: Optional[Sequence[str]] = None,
    skip_empty: bool = True,
    min_abs_weight: int = 1,
) -> List[str]:
    """``IF ... THEN`` rules listing every output the clause votes for/against with its weight."""
    names = list(class_names or self.class_names or [str(k) for k in range(self.n_outputs)])
    out = []
    W = self.weights.tolist()
    for j in range(self.n_clauses_total):
        if skip_empty and int(self.include_count[j]) == 0:
            continue
        votes = []
        for i, w in enumerate(W[j]):
            if abs(w) >= min_abs_weight:
                votes.append(f"{'+' if w >= 0 else '-'}{abs(w)} {names[i]}")
        if not votes:
            continue
        out.append(f"IF {self.clause_expression(j, feature_names)} THEN {', '.join(votes)}")
    return out

torchtsetlin.models.RegressionTsetlinMachine

RegressionTsetlinMachine(n_features: Optional[int], n_clauses: int, T: float, s: float = 10.0, *, y_range: Optional[Tuple[float, float]] = None, weighted: bool = False, max_weight: Optional[int] = None, **kwargs)

Bases: TsetlinMachineBase

Regression Tsetlin machine (Abeyrathna et al.).

All clauses have positive polarity; the vote sum v = sum_j w_j c_j is clamped to [0, T] and mapped linearly onto the target range: y_hat = y_min + (v / T) * (y_max - y_min).

Learning: with the target expressed in vote units v* = T (y - y_min)/(y_max - y_min), every clause receives Type I feedback with probability ((v* - v)/T)**2 when the prediction is too low, and Type II feedback with probability ((v - v*)/T)**2 when it is too high. The squared error is the rule of the reference implementations (pyTsetlinMachine, pyTsetlinMachineParallel and tmu).

Parameters:

Name	Type	Description	Default
n_features	Optional[int]	Boolean feature count (None = lazy).	required
n_clauses	int	number of clauses.	required
T	float	vote margin — also the output resolution (T distinct predictions).	required
s	float	specificity.	10.0
y_range	Optional[Tuple[float, float]]	(y_min, y_max) of the targets. If None it is inferred from the first training batch (a warning is issued); prefer to pass it explicitly.	None
weighted	bool	learn integer clause weights (Type Ia +1, Type II -1, min 1).	False
max_weight	Optional[int]	optional upper bound on the clause weights (None: unbounded).	None
**kwargs	Any	forwarded to :class:~torchtsetlin.models.TsetlinMachineBase.	{}

Source code in src/torchtsetlin/models/regression.py

def __init__(
    self,
    n_features: Optional[int],
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    y_range: Optional[Tuple[float, float]] = None,
    weighted: bool = False,
    max_weight: Optional[int] = None,
    **kwargs,
) -> None:
    if T <= 0:
        raise ValueError("T must be positive")
    self.n_clauses = int(n_clauses)
    self.T = float(T)
    self.weighted = bool(weighted)
    if max_weight is not None and int(max_weight) < 1:
        raise ValueError("max_weight must be >= 1")
    self.max_weight = None if max_weight is None else int(max_weight)
    self.y_min: Optional[float] = None
    self.y_max: Optional[float] = None
    if y_range is not None:
        self.set_range(*y_range)
    super().__init__(n_features, self.n_clauses, s=s, **kwargs)
    self.register_buffer("weights", torch.ones(self.n_clauses, dtype=torch.int32))

predict

predict(x) -> Tensor

Continuous predictions (B,) in [y_min, y_max].

Source code in src/torchtsetlin/models/regression.py

def predict(self, x) -> Tensor:
    """Continuous predictions ``(B,)`` in ``[y_min, y_max]``."""
    return self.votes_to_targets(self.forward(x))

rules

rules(feature_names: Optional[Sequence[str]] = None, skip_empty: bool = True) -> List[str]

IF ... THEN +w rules: each matching clause adds w * (y_max - y_min) / T.

Source code in src/torchtsetlin/models/regression.py

def rules(self, feature_names: Optional[Sequence[str]] = None, skip_empty: bool = True) -> List[str]:
    """``IF ... THEN +w`` rules: each matching clause adds ``w * (y_max - y_min) / T``."""
    self._check_range()
    step = (self.y_max - self.y_min) / self.T  # type: ignore[operator]
    out = []
    for j in range(self.n_clauses_total):
        if skip_empty and int(self.include_count[j]) == 0:
            continue
        w = int(self.weights[j]) if self.weighted else 1
        out.append(f"IF {self.clause_expression(j, feature_names)} THEN +{w * step:.4g}")
    return out

torchtsetlin.models.ConvTsetlinMachine

ConvTsetlinMachine(n_classes: int, n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrPair = 10, stride: IntOrPair = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int, int]] = None, **kwargs)

Bases: _ConvMixin, TsetlinMachine

Convolutional multi-class Tsetlin machine (Granmo et al., 2019).

Parameters:

Name	Type	Description	Default
n_classes	int	number of classes.	required
n_clauses	int	clauses per class.	required
T	float	vote margin.	required
s	float	specificity.	10.0
patch_size	IntOrPair	(kh, kw) convolution window (int for square windows).	10
stride	IntOrPair	window stride.	1
position_encoding	bool	append thermometer-encoded patch coordinates to each patch.	True
input_shape	Optional[Tuple[int, int, int]]	(Z, H, W) of the Boolean images; None = infer from the first batch.	None
**kwargs	Any	forwarded to :class:~torchtsetlin.models.TsetlinMachine (weighted, n_states, max_included_literals, drop_clause_p, ...).	{}

Shape

• input: (B, Z, H, W) Boolean ((B, H, W) is treated as one plane).
• output: (B, n_classes) vote sums.

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_classes: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrPair = 10,
    stride: IntOrPair = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int, int]] = None,
    **kwargs,
) -> None:
    self._init_conv(patch_size, stride, position_encoding, input_shape)
    n_features = self.conv_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_classes, n_clauses, T, s, **kwargs)

torchtsetlin.models.ConvCoalescedTsetlinMachine

ConvCoalescedTsetlinMachine(n_outputs: int, n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrPair = 10, stride: IntOrPair = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int, int]] = None, **kwargs)

Bases: _ConvMixin, CoalescedTsetlinMachine

Convolutional coalesced Tsetlin machine (shared clauses, per-output weights).

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_outputs: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrPair = 10,
    stride: IntOrPair = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int, int]] = None,
    **kwargs,
) -> None:
    self._init_conv(patch_size, stride, position_encoding, input_shape)
    n_features = self.conv_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_outputs, n_clauses, T, s, **kwargs)

torchtsetlin.models.ConvRegressionTsetlinMachine

ConvRegressionTsetlinMachine(n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrPair = 10, stride: IntOrPair = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int, int]] = None, **kwargs)

Bases: _ConvMixin, RegressionTsetlinMachine

Convolutional regression Tsetlin machine (Abeyrathna et al., 2021).

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrPair = 10,
    stride: IntOrPair = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int, int]] = None,
    **kwargs,
) -> None:
    self._init_conv(patch_size, stride, position_encoding, input_shape)
    n_features = self.conv_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_clauses, T, s, **kwargs)

torchtsetlin.models.Conv1dTsetlinMachine

Conv1dTsetlinMachine(n_classes: int, n_clauses: int, T: float, s: float = 10.0, *, kernel_size: int = 5, stride: int = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int]] = None, **kwargs)

Bases: ConvTsetlinMachine

1D convolutional Tsetlin machine for sequences / time series.

Input (B, Z, L) (Z Boolean channels, length L) is handled as a (B, Z, 1, L) image with a 1 x kernel_size window.

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_classes: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    kernel_size: int = 5,
    stride: int = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int]] = None,
    **kwargs,
) -> None:
    shape3 = (input_shape[0], 1, input_shape[1]) if input_shape is not None else None
    super().__init__(
        n_classes,
        n_clauses,
        T,
        s,
        patch_size=(1, kernel_size),
        stride=(1, stride),
        position_encoding=position_encoding,
        input_shape=shape3,
        **kwargs,
    )

torchtsetlin.models.Conv3dTsetlinMachine

Conv3dTsetlinMachine(n_classes: int, n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrTriple = (3, 7, 7), stride: IntOrTriple = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int, int, int]] = None, patches_per_chunk: Optional[int] = None, **kwargs)

Bases: _Conv3dMixin, TsetlinMachine

Convolutional multi-class Tsetlin machine for volumes (3D windows).

The 3D generalisation of :class:ConvTsetlinMachine: windows slide over depth, height and width, so one clause can find a pattern anywhere in a volume — an MRI/CT stack, or a (time, H, W) sequence of frames.

Parameters:

Name	Type	Description	Default
n_classes	int	number of classes.	required
n_clauses	int	clauses per class.	required
T	float	vote margin.	required
s	float	specificity.	10.0
patch_size	IntOrTriple	(kd, kh, kw) window (int for a cube).	(3, 7, 7)
stride	IntOrTriple	window stride (int or (sd, sh, sw)).	1
position_encoding	bool	append thermometer-encoded patch coordinates (depth, row, column) to each patch.	True
input_shape	Optional[Tuple[int, int, int, int]]	(Z, D, H, W) of the Boolean volumes; None = infer from the first batch.	None
patches_per_chunk	Optional[int]	evaluate at most this many patches together (bounds memory on large volumes). None splits the patch axis only when one example would exceed max_chunk_elements.	None
**kwargs	Any	forwarded to :class:~torchtsetlin.models.TsetlinMachine.	{}

Shape

• input: (B, Z, D, H, W) Boolean ((B, D, H, W) is treated as one plane).
• output: (B, n_classes) vote sums.

Example

model = Conv3dTsetlinMachine(2, 40, T=15, patch_size=3, position_encoding=False) model.update(volumes, labels) # volumes: (B, 1, 8, 16, 16) Boolean model.clause_patch(0).shape # (Z, kd, kh, kw)

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_classes: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrTriple = (3, 7, 7),
    stride: IntOrTriple = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int, int, int]] = None,
    patches_per_chunk: Optional[int] = None,
    **kwargs,
) -> None:
    self._init_conv3d(patch_size, stride, position_encoding, input_shape, patches_per_chunk)
    n_features = self.conv_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_classes, n_clauses, T, s, **kwargs)

torchtsetlin.models.Conv3dCoalescedTsetlinMachine

Conv3dCoalescedTsetlinMachine(n_outputs: int, n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrTriple = (3, 7, 7), stride: IntOrTriple = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int, int, int]] = None, patches_per_chunk: Optional[int] = None, **kwargs)

Bases: _Conv3dMixin, CoalescedTsetlinMachine

3D convolutional coalesced Tsetlin machine (shared clauses, per-output weights).

Arguments as :class:Conv3dTsetlinMachine, with n_outputs instead of n_classes and n_clauses counting the shared clause pool.

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_outputs: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrTriple = (3, 7, 7),
    stride: IntOrTriple = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int, int, int]] = None,
    patches_per_chunk: Optional[int] = None,
    **kwargs,
) -> None:
    self._init_conv3d(patch_size, stride, position_encoding, input_shape, patches_per_chunk)
    n_features = self.conv_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_outputs, n_clauses, T, s, **kwargs)

torchtsetlin.models.Conv3dRegressionTsetlinMachine

Conv3dRegressionTsetlinMachine(n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrTriple = (3, 7, 7), stride: IntOrTriple = 1, position_encoding: bool = True, input_shape: Optional[Tuple[int, int, int, int]] = None, patches_per_chunk: Optional[int] = None, **kwargs)

Bases: _Conv3dMixin, RegressionTsetlinMachine

3D convolutional regression Tsetlin machine.

Arguments as :class:Conv3dTsetlinMachine, without n_classes.

Source code in src/torchtsetlin/models/conv.py

def __init__(
    self,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrTriple = (3, 7, 7),
    stride: IntOrTriple = 1,
    position_encoding: bool = True,
    input_shape: Optional[Tuple[int, int, int, int]] = None,
    patches_per_chunk: Optional[int] = None,
    **kwargs,
) -> None:
    self._init_conv3d(patch_size, stride, position_encoding, input_shape, patches_per_chunk)
    n_features = self.conv_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_clauses, T, s, **kwargs)

torchtsetlin.models.SegmentationTsetlinMachine

SegmentationTsetlinMachine(n_classes: int, n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrPair = 3, stride: IntOrPair = 1, padding: Union[str, IntOrPair] = 'same', position_encoding: PositionEncoding = False, input_shape: Optional[Tuple[int, int, int]] = None, patches_per_commit: Optional[int] = 128, ignore_index: Optional[int] = None, pixel_order: str = 'raster', class_feedback_p: Optional[Tensor] = None, **kwargs)

Bases: _DenseMixin, TsetlinMachine

Dense per-pixel Tsetlin machine with class-owned clauses.

Each class owns n_clauses clauses (half voting for it, half against), exactly as in :class:~torchtsetlin.models.TsetlinMachine; the difference is that a pixel, not an image, is the example. Predictions are label maps.

Parameters:

Name	Type	Description	Default
n_classes	int	number of segmentation classes.	required
n_clauses	int	clauses per class.	required
T	float	vote margin.	required
s	float	specificity.	10.0
patch_size	IntOrPair	(kh, kw) neighbourhood each pixel is classified from.	3
stride	IntOrPair	output-grid stride (1 keeps the input resolution).	1
padding	Union[str, IntOrPair]	"same" (default) or an explicit (ph, pw) pair.	'same'
position_encoding	PositionEncoding	False, True (one thermometer bit per output row/column) or an int giving the number of bits per axis. Position literals let clauses specialise by location — on scenes with a stable layout (a horizon, a scanner field of view) this matters a lot; see docs/concepts/segmentation.md.	False
input_shape	Optional[Tuple[int, int, int]]	(Z, H, W) of the Boolean images; None = infer from the first batch.	None
patches_per_commit	Optional[int]	commit accumulated feedback every this many pixels (default 128). None commits once per image chunk — faster, and less faithful to the classical algorithm the larger the images are.	128
pixel_order	str	"raster" (default) commits pixels in image-major raster order; "shuffle" permutes all pixels of an image chunk first, so each commit mixes pixels from different images and positions. With large commits this is markedly better than raster order (segmentation experiments, H1); the chunk holds as many images as max_chunk_elements allows for the encoded literals.	'raster'
ignore_index	Optional[int]	label value that contributes no feedback and is excluded from metrics (the "void" class of CamVid / Cityscapes / Pascal VOC).	None
class_feedback_p	Optional[Tensor]	optional (n_classes,) probability that a pixel of each class takes part in feedback — the imbalance lever for dense tasks, since a Tsetlin machine has no loss to weight. See :func:torchtsetlin.data.balanced_class_probabilities.	None
**kwargs	Any	forwarded to :class:~torchtsetlin.models.TsetlinMachine (weighted, n_states, max_included_literals, drop_clause_p, ...).	{}

Shape

• input: (B, Z, H, W) Boolean ((B, H, W) is treated as one plane).
• :meth:forward: (B*Py*Px, n_classes) vote sums (patches folded into the batch).
• :meth:vote_map: (B, n_classes, Py, Px); :meth:predict: (B, Py, Px).

Example

model = SegmentationTsetlinMachine(4, n_clauses=200, T=60, s=10.0, patch_size=3) model.update(x, labels) # x: (B, Z, H, W) bool, labels: (B, H, W) long model.eval(); pred = model.predict(x) # (B, H, W)

Source code in src/torchtsetlin/models/segmentation.py

def __init__(
    self,
    n_classes: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrPair = 3,
    stride: IntOrPair = 1,
    padding: Union[str, IntOrPair] = "same",
    position_encoding: PositionEncoding = False,
    input_shape: Optional[Tuple[int, int, int]] = None,
    patches_per_commit: Optional[int] = 128,
    ignore_index: Optional[int] = None,
    pixel_order: str = "raster",
    class_feedback_p: Optional[Tensor] = None,
    **kwargs,
) -> None:
    self._init_dense(
        patch_size, stride, padding, position_encoding, input_shape, patches_per_commit,
        ignore_index, pixel_order,
    )
    n_features = self.dense_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_classes, n_clauses, T, s, **kwargs)
    # A (non-persistent) buffer so ``.to(device)`` moves it with the automata; a plain
    # tensor attribute stayed on the construction device and crashed GPU training.
    self.register_buffer("class_feedback_p", None, persistent=False)
    self.set_class_feedback_p(class_feedback_p)

set_class_feedback_p

set_class_feedback_p(p: Optional[Tensor]) -> None

Set (or clear) the per-class probability that a pixel contributes feedback.

Source code in src/torchtsetlin/models/segmentation.py

def set_class_feedback_p(self, p: Optional[Tensor]) -> None:
    """Set (or clear) the per-class probability that a pixel contributes feedback."""
    if p is None:
        self.class_feedback_p = None
        return
    t = torch.as_tensor(p, dtype=torch.float32, device=self.ta_state.device).view(-1)
    if t.numel() != self.n_classes:
        raise ValueError(f"class_feedback_p must have {self.n_classes} entries")
    if bool(((t < 0) | (t > 1)).any()):
        raise ValueError("class_feedback_p entries must lie in [0, 1]")
    self.class_feedback_p = t

torchtsetlin.models.CoalescedSegmentationTsetlinMachine

CoalescedSegmentationTsetlinMachine(n_outputs: int, n_clauses: int, T: float, s: float = 10.0, *, patch_size: IntOrPair = 3, stride: IntOrPair = 1, padding: Union[str, IntOrPair] = 'same', position_encoding: PositionEncoding = False, input_shape: Optional[Tuple[int, int, int]] = None, patches_per_commit: Optional[int] = 128, ignore_index: Optional[int] = None, pixel_order: str = 'raster', **kwargs)

Bases: _DenseMixin, CoalescedTsetlinMachine

Dense per-pixel Tsetlin machine with one shared clause pool and per-class weights.

The coalesced output layer (Glimsdal & Granmo, 2021) is a better fit than class-owned clauses when classes share structure — edges, textures and gradients are the same features whichever class they belong to — because one clause can serve every class it is useful for instead of being relearned per class.

With multi_label=True the targets are a Boolean stack of masks (B, K, H, W) and a pixel may belong to several classes at once (overlapping structures); ignore_index does not apply in that mode.

Parameters:

Name	Type	Description	Default
n_outputs	int	number of classes (or of overlapping masks in multi-label mode).	required
n_clauses	int	total number of shared clauses.	required
T	float	vote margin.	required
s	float	specificity.	10.0
**kwargs	Any	the dense arguments of :class:~torchtsetlin.models.SegmentationTsetlinMachine plus everything :class:~torchtsetlin.models.CoalescedTsetlinMachine accepts.	{}

Source code in src/torchtsetlin/models/segmentation.py

def __init__(
    self,
    n_outputs: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    patch_size: IntOrPair = 3,
    stride: IntOrPair = 1,
    padding: Union[str, IntOrPair] = "same",
    position_encoding: PositionEncoding = False,
    input_shape: Optional[Tuple[int, int, int]] = None,
    patches_per_commit: Optional[int] = 128,
    ignore_index: Optional[int] = None,
    pixel_order: str = "raster",
    **kwargs,
) -> None:
    self._init_dense(
        patch_size, stride, padding, position_encoding, input_shape, patches_per_commit,
        ignore_index, pixel_order,
    )
    n_features = self.dense_features(self.input_shape) if self.input_shape else None
    super().__init__(n_features, n_outputs, n_clauses, T, s, **kwargs)
    if self.multi_label and self.ignore_index is not None:
        raise ValueError("ignore_index is not supported in multi-label mode")

predict

predict(x) -> Tensor

Label map (B, Py, Px), or a Boolean mask stack (B, K, Py, Px) in multi-label mode.

Source code in src/torchtsetlin/models/segmentation.py

@torch.no_grad()
def predict(self, x) -> Tensor:  # type: ignore[override]
    """Label map ``(B, Py, Px)``, or a Boolean mask stack ``(B, K, Py, Px)`` in
    multi-label mode."""
    votes = self.vote_map(x)
    return votes > 0 if self.multi_label else votes.argmax(dim=1)

torchtsetlin.models.GraphTsetlinMachine

GraphTsetlinMachine(n_features: Optional[int], n_outputs: int, n_clauses: int, T: float, s: float = 10.0, *, depth: int = 1, message_size: int = 256, message_bits: int = 2, message_encoding: str = 'random', n_edge_types: int = 1, n_node_types: int = 1, shared_node_draw: bool = False, edge_type_names: Optional[Sequence[str]] = None, **kwargs)

Bases: _GraphMixin, CoalescedTsetlinMachine

Graph Tsetlin machine (Granmo et al., 2025): deep clauses over graph-structured input.

Every clause has depth components. Component 0 matches a node's own Boolean features; component i >= 1 matches the node's inbox of layer i, which holds a message (clause k, edge type t) for every in-neighbour, reached over an edge of type t, at which clause k's components 0 .. i-1 all match. A clause is True for a graph when its components 0 .. depth-1 all match at one node. Votes, clause sharing and the feedback probabilities are those of :class:CoalescedTsetlinMachine; each feedback event updates every layer of the clause from one random node where the whole clause matches.

Input: a :class:~torchtsetlin.data.GraphBatch. With depth=1 and the graphs of :meth:GraphBatch.from_images, this is exactly :class:~torchtsetlin.models.ConvCoalescedTsetlinMachine.

Parameters:

Name	Type	Description	Default
n_features	Optional[int]	Boolean features per node F (None = lazy, from the first batch).	required
n_outputs	int	number of classes (multi-class) or labels (multi_label=True).	required
n_clauses	int	number of shared clauses.	required
T	float	vote margin.	required
s	float	specificity.	10.0
depth	int	number of layers D (1 = no message passing; each extra layer reaches one hop further).	1
message_size	int	inbox size M per layer for hashed messages. One-hot messages use M = n_clauses * n_edge_types and ignore it.	256
message_bits	int	positions per clause message (hashed "random" encoding).	2
message_encoding	str	"random" (default, as the reference implementation), "double_hashing" (2 positions: j mod M and p - (j mod p)) or "one_hot" (one position per clause and edge type: exact, collision free, and readable as the paper's r1:0 notation, but the inbox grows with the clause count).	'random'
n_edge_types	int	edge types of the dataset (a message is bound to its edge type by a cyclic shift of its positions).	1
n_node_types	int	node types; clause j only matches nodes of type j mod n_node_types, as in the reference implementation.	1
shared_node_draw	bool	draw one random matching node per (graph, clause) and use it for all classes and for Type I and II, as the reference implementation does. The default draws a node per feedback event, as the convolutional models do.	False
edge_type_names	Optional[Sequence[str]]	names used when printing message literals (default e0, e1, ...).	None
**kwargs		forwarded to :class:CoalescedTsetlinMachine (multi_label, negative_scale, max_included_literals, feedback_mode, ...).	{}

The reference implementation updates every class with probability scaled by q / (K - 1) for the negative ones; multi_label=True with one-hot targets and negative_scale = q / (K - 1) is that rule. It also learns one graph at a time: feedback_mode="sequential".

Source code in src/torchtsetlin/models/graph.py

def __init__(
    self,
    n_features: Optional[int],
    n_outputs: int,
    n_clauses: int,
    T: float,
    s: float = 10.0,
    *,
    depth: int = 1,
    message_size: int = 256,
    message_bits: int = 2,
    message_encoding: str = "random",
    n_edge_types: int = 1,
    n_node_types: int = 1,
    shared_node_draw: bool = False,
    edge_type_names: Optional[Sequence[str]] = None,
    **kwargs,
) -> None:
    if depth < 1:
        raise ValueError("depth must be >= 1")
    if message_encoding not in _ENCODINGS:
        raise ValueError(f"message_encoding must be one of {_ENCODINGS}")
    if n_edge_types < 1 or n_node_types < 1:
        raise ValueError("n_edge_types and n_node_types must be >= 1")
    C = int(n_clauses)
    if message_encoding == "one_hot":
        message_size, message_bits = C * int(n_edge_types), 1
    elif message_encoding == "double_hashing":
        message_bits = 2
    if message_encoding != "one_hot" and not 1 <= message_bits <= message_size:
        raise ValueError("message_bits must lie in [1, message_size]")
    self.depth = int(depth)
    self.message_size = int(message_size)
    self.message_bits = int(message_bits)
    self.message_encoding = message_encoding
    self.n_edge_types = int(n_edge_types)
    self.n_node_types = int(n_node_types)
    self.shared_node_draw = bool(shared_node_draw)
    self.edge_type_names = None if edge_type_names is None else [str(n) for n in edge_type_names]
    if self.edge_type_names is not None and len(self.edge_type_names) != self.n_edge_types:
        raise ValueError("edge_type_names needs one name per edge type")
    total = None if n_features is None else int(n_features) + (self.depth - 1) * self.message_size
    super().__init__(total, n_outputs, n_clauses, T, s, **kwargs)
    self.register_buffer("msg_bits", self._make_message_bits())

torchtsetlin.models.TsetlinMachineBase

TsetlinMachineBase(n_features: Optional[int], n_clauses_total: int, *, n_states: int = 128, s: float = 10.0, boost_true_positive: bool = True, max_included_literals: Optional[int] = None, drop_clause_p: float = 0.0, drop_literal_p: float = 0.0, drop_granularity: str = 'batch', init: str = 'boundary', state_dtype: Optional[dtype] = None, compute_dtype: dtype = torch.float32, max_chunk_elements: int = 2 ** 27, feedback_mode: str = 'batch', budget_mode: str = 'soft', clamp_prediction_votes: Optional[bool] = None)

Bases: Module

Common machinery for Tsetlin machines: automata state, clause evaluation, dropout, and the batched learning loop.

The learnable state lives in the integer buffer ta_state of shape (n_clauses_total, 2 * n_features). It moves with .to(device), is saved in state_dict() and can be inspected directly.

Subclasses implement the output layer (how clause outputs become votes) and the feedback policy (which clauses receive Type I / Type II feedback for a labelled example) through :meth:_votes, :meth:_select_feedback and the weight hooks. Convolutional variants additionally override :meth:_encode, :meth:_evaluate and :meth:_feedback_counts.

Parameters:

Name	Type	Description	Default
n_features	Optional[int]	number of Boolean input features F. None defers allocation until the first batch is seen (lazy initialisation, like torch.nn.LazyLinear).	required
n_clauses_total	int	total number of clauses C (all classes / outputs together).	required
n_states	int	automaton states per action N. A literal is included when its state is >= N. Larger values give a deeper memory (slower forgetting).	128
s	float	specificity. Memorize probability (s-1)/s, forget probability 1/s. Larger s -> more literals per clause (more specific patterns).	10.0
boost_true_positive	bool	memorize True literals of matching clauses with probability 1 instead of (s-1)/s (the default in the reference implementations).	True
max_included_literals	Optional[int]	optional budget on the number of included literals per clause (clause size constraint). A matching clause above the budget receives only forgetting (Type Ib) feedback until it shrinks back (see budget_mode). Clause weights are still updated for such a clause, as in tmu (pyTsetlinMachine skips the weight increment of an over-budget clause).	None
drop_clause_p	float	probability of dropping a clause during learning (a Tsetlin machine analogue of dropout). The mask is resampled on every :meth:update call, or once per epoch if drop_granularity="epoch" (call :meth:resample_dropout).	0.0
drop_literal_p	float	probability of freezing a literal (no feedback) per resample.	0.0
drop_granularity	str	"batch" or "epoch".	'batch'
init	str	"boundary" starts every automaton just below the include boundary (N-1); "random" starts randomly at N-1 or N.	'boundary'
clamp_prediction_votes	Optional[bool]	clamp vote sums to [-T, T] in forward() / predict() (pyTsetlinMachine's rule, under which saturated classes tie and the lowest index wins). None keeps the model's default (unclamped for classifiers, as in tmu).	None
state_dtype	Optional[dtype]	integer dtype of the automata states. None (default) picks the smallest safe one: uint8 when 2 * n_states - 1 <= 255 (8 bits per automaton, as in the reference implementations), else int16 or int32. Cast before doing arithmetic on ta_state yourself (model.ta_state.long() - N).	None
compute_dtype	dtype	floating dtype of the literals and feedback counts. On CUDA, clauses are evaluated in float16 regardless: the evaluation only tests whether a count of violated literals is zero, which half precision decides exactly (a sum of non-negative integers never rounds to zero).	float32
max_chunk_elements	int	memory budget (in tensor elements) for the largest intermediate tensor of an update; batches are processed in chunks that respect it.	2 ** 27
feedback_mode	str	"batch" accumulates the feedback of all examples in a batch and applies it once (fast, mini-batch style). "sequential" applies the feedback example by example, re-evaluating the clauses in between — the exact classical algorithm, at the cost of one small step per example.	'batch'
budget_mode	str	how max_included_literals is enforced. "soft" (default) is the reference implementation's rule (tmu): the budget only gates Type Ia, so a clause can overshoot it and is then pruned by Type Ib, which keeps clause contents changing. "hard" (the 0.3.x default) also reverts every inclusion beyond the budget. A clause then never exceeds it, so the Type Ib switch never fires and clauses freeze at exactly the budget, which cost 10-11 pp on CIFAR-10 colour-thermometer models.	'soft'

Source code in src/torchtsetlin/models/base.py

def __init__(
    self,
    n_features: Optional[int],
    n_clauses_total: int,
    *,
    n_states: int = 128,
    s: float = 10.0,
    boost_true_positive: bool = True,
    max_included_literals: Optional[int] = None,
    drop_clause_p: float = 0.0,
    drop_literal_p: float = 0.0,
    drop_granularity: str = "batch",
    init: str = "boundary",
    state_dtype: Optional[torch.dtype] = None,
    compute_dtype: torch.dtype = torch.float32,
    max_chunk_elements: int = 2**27,
    feedback_mode: str = "batch",
    budget_mode: str = "soft",
    clamp_prediction_votes: Optional[bool] = None,
) -> None:
    super().__init__()
    if clamp_prediction_votes is not None:
        # instance override of the class default (pyTsetlinMachine / PyTsetlinMachineCUDA
        # take the argmax over votes clamped to [-T, T]; tmu and green-tsetlin do not)
        self.clamp_prediction_votes = bool(clamp_prediction_votes)
    if n_clauses_total <= 0:
        raise ValueError("n_clauses_total must be positive")
    if n_states < 2:
        raise ValueError("n_states must be >= 2")
    if s < 1.0:
        raise ValueError("s (specificity) must be >= 1")
    if not (0.0 <= drop_clause_p < 1.0) or not (0.0 <= drop_literal_p < 1.0):
        raise ValueError("dropout probabilities must lie in [0, 1)")
    if init not in ("boundary", "random"):
        raise ValueError("init must be 'boundary' or 'random'")
    if drop_granularity not in ("batch", "epoch"):
        raise ValueError("drop_granularity must be 'batch' or 'epoch'")
    if feedback_mode not in ("batch", "sequential"):
        raise ValueError("feedback_mode must be 'batch' or 'sequential'")
    if budget_mode not in ("hard", "soft"):
        raise ValueError("budget_mode must be 'hard' or 'soft'")
    self.n_features: Optional[int] = None
    self.n_clauses_total = int(n_clauses_total)
    self.n_states = int(n_states)
    self.s = float(s)
    self.boost_true_positive = bool(boost_true_positive)
    self.max_included_literals = max_included_literals
    self.drop_clause_p = float(drop_clause_p)
    self.drop_literal_p = float(drop_literal_p)
    self.drop_granularity = drop_granularity
    self.init = init
    if state_dtype is None:
        # 8 bits per automaton, as in the reference implementations, when 2N - 1 fits.
        top = 2 * int(n_states) - 1
        state_dtype = torch.uint8 if top <= 255 else torch.int16 if top <= 32767 else torch.int32
    self.state_dtype = state_dtype
    self.compute_dtype = compute_dtype
    self.max_chunk_elements = int(max_chunk_elements)
    self.feedback_mode = feedback_mode
    self.budget_mode = budget_mode
    self.feature_names: Optional[List[str]] = None
    self._clause_active: Optional[Tensor] = None
    self._literal_active: Optional[Tensor] = None

    # Buffers (empty until the feature count is known).
    self.register_buffer("ta_state", torch.empty(0, dtype=self.state_dtype))
    self.register_buffer("include", torch.empty(0, dtype=compute_dtype), persistent=False)
    self.register_buffer("include_count", torch.empty(0, dtype=torch.long), persistent=False)
    self.register_load_state_dict_post_hook(_refresh_after_load)
    if n_features is not None:
        self._initialize_state(int(n_features))

reset_parameters

reset_parameters() -> None

Re-initialise all automata (keeps the shape).

Source code in src/torchtsetlin/models/base.py

def reset_parameters(self) -> None:
    """Re-initialise all automata (keeps the shape)."""
    if self.is_initialized:
        self._initialize_state(int(self.n_features), self.ta_state.device)  # type: ignore[arg-type]

clause_outputs

clause_outputs(literals: Tensor, empty_value: Optional[bool] = None) -> Tensor

Evaluate every clause on encoded literals -> (B, C) bool.

empty_value defaults to self.training (empty clauses are True while learning and False when predicting).

Source code in src/torchtsetlin/models/base.py

def clause_outputs(self, literals: Tensor, empty_value: Optional[bool] = None) -> Tensor:
    """Evaluate every clause on encoded literals -> ``(B, C)`` bool.

    ``empty_value`` defaults to ``self.training`` (empty clauses are True while learning
    and False when predicting).
    """
    if empty_value is None:
        empty_value = self.training
    return self._evaluate(literals, empty_value)[0]

evaluate_clauses

evaluate_clauses(x) -> Tensor

Clause outputs for raw Boolean input x (prediction semantics) -> (B, C).

Source code in src/torchtsetlin/models/base.py

def evaluate_clauses(self, x) -> Tensor:
    """Clause outputs for raw Boolean input ``x`` (prediction semantics) -> ``(B, C)``."""
    with torch.no_grad():
        xb = self._prepare(x)
        outs = []
        for sl in chunk_indices(xb.shape[0], self._chunk_size(xb)):
            outs.append(self._evaluate(self._encode(xb[sl]), empty_value=False)[0])
        return torch.cat(outs, dim=0)

forward

forward(x) -> Tensor

Vote sums (B, n_outputs). Uses prediction semantics for empty clauses when the module is in eval() mode. The sums are not clamped to [-T, T] (unless the model sets clamp_prediction_votes), as in the reference implementations.

Source code in src/torchtsetlin/models/base.py

def forward(self, x) -> Tensor:
    """Vote sums ``(B, n_outputs)``. Uses prediction semantics for empty clauses when the
    module is in ``eval()`` mode. The sums are *not* clamped to ``[-T, T]`` (unless the
    model sets ``clamp_prediction_votes``), as in the reference implementations."""
    with torch.no_grad():
        xb = self._prepare(x)
        votes = []
        for sl in chunk_indices(xb.shape[0], self._chunk_size(xb)):
            out, _ = self._evaluate(self._encode(xb[sl]), empty_value=self.training)
            votes.append(self._votes(out, clamp=self.clamp_prediction_votes))
        return torch.cat(votes, dim=0)

update

update(x, y, sequential: Optional[bool] = None) -> Tensor

One learning step on a batch — the Tsetlin machine analogue of loss.backward(); optimizer.step().

In "batch" mode every example of the batch is evaluated against the current automata; the resulting Type I / Type II feedback events are accumulated and applied once. With a batch size of 1 this is exactly the classical sequential algorithm; larger batches trade a little fidelity for throughput (like mini-batch SGD). In "sequential" mode (feedback_mode or the sequential argument) the examples are processed one by one with the automata updated in between.

Parameters:

Name	Type	Description	Default
x	Tensor | ndarray | list	Boolean input batch (any Boolean-like dtype).	required
y	Tensor | ndarray | list	targets (class indices, multi-label matrix or regression values, depending on the model).	required
sequential	Optional[bool]	override self.feedback_mode for this call.	None

Returns:

Type	Description
Tensor	The vote sums computed before the update ((B, n_outputs)) so that batch
Tensor	metrics can be logged without a second forward pass.

Source code in src/torchtsetlin/models/base.py

def update(self, x, y, sequential: Optional[bool] = None) -> Tensor:
    """One learning step on a batch — the Tsetlin machine analogue of
    ``loss.backward(); optimizer.step()``.

    In ``"batch"`` mode every example of the batch is evaluated against the current
    automata; the resulting Type I / Type II feedback events are accumulated and applied
    once. With a batch size of 1 this is exactly the classical sequential algorithm;
    larger batches trade a little fidelity for throughput (like mini-batch SGD). In
    ``"sequential"`` mode (``feedback_mode`` or the ``sequential`` argument) the examples
    are processed one by one with the automata updated in between.

    Args:
        x (Tensor | ndarray | list): Boolean input batch (any Boolean-like dtype).
        y (Tensor | ndarray | list): targets (class indices, multi-label matrix or
            regression values, depending on the model).
        sequential: override ``self.feedback_mode`` for this call.

    Returns:
        The vote sums computed **before** the update (``(B, n_outputs)``) so that batch
        metrics can be logged without a second forward pass.
    """
    with torch.no_grad():
        xb = self._prepare(x)
        B = xb.shape[0]
        y_t = self._coerce_targets(y, B, xb.device)
        if self.drop_granularity == "batch":
            self.resample_dropout()
        seq = self.feedback_mode == "sequential" if sequential is None else bool(sequential)
        # sequential mode: whether to fuse is decided on the work of the whole call
        fuse = seq and self._fuse_single_examples(B)
        if self._graphs_ok(seq, xb, fuse):
            return self._graphed_update(xb, y_t, seq, fuse)
        return self._update_block(xb, y_t, seq, fuse)

train

train(mode: bool = True) -> TsetlinMachineBase

As :meth:torch.nn.Module.train; switching to eval() also releases the count buffers that learning keeps between updates.

Source code in src/torchtsetlin/models/base.py

def train(self, mode: bool = True) -> TsetlinMachineBase:
    """As :meth:`torch.nn.Module.train`; switching to ``eval()`` also releases the count
    buffers that learning keeps between updates."""
    if not mode:
        self.__dict__.pop("_acc_cache", None)
    return super().train(mode)

include_bits

include_bits() -> Tensor

The include mask bit-packed, (C, ceil(2F/32)) int32 (literal l in bit l % 32 of word l // 32), refreshed from :attr:include when stale. The Triton commit writes it directly.

Source code in src/torchtsetlin/models/base.py

def include_bits(self) -> Tensor:
    """The include mask bit-packed, ``(C, ceil(2F/32))`` int32 (literal ``l`` in bit
    ``l % 32`` of word ``l // 32``), refreshed from :attr:`include` when stale. The Triton
    commit writes it directly."""
    bits = self.include_bits_buffer()
    if not self.__dict__.get("_bits_fresh", False):
        bits.copy_(_triton.pack_bits(self.include))
        self.__dict__["_bits_fresh"] = True
    return bits

include_bits_buffer

include_bits_buffer() -> Tensor

The (C, ceil(2F/32)) int32 buffer behind :meth:include_bits, without refreshing it (a commit kernel overwrites it whole).

Source code in src/torchtsetlin/models/base.py

def include_bits_buffer(self) -> Tensor:
    """The ``(C, ceil(2F/32))`` int32 buffer behind :meth:`include_bits`, without refreshing
    it (a commit kernel overwrites it whole)."""
    bits = self.__dict__.get("_include_bits")
    W = _triton.n_words(self.n_literals)
    if (bits is None or bits.device != self.ta_state.device
            or tuple(bits.shape) != (self.n_clauses_total, W)):
        bits = torch.zeros(self.n_clauses_total, W, dtype=torch.int32, device=self.ta_state.device)
        self.__dict__["_include_bits"] = bits
        self.__dict__["_bits_fresh"] = False
    return bits

resample_dropout

resample_dropout() -> None

Draw new clause / literal dropout masks (called automatically per update when drop_granularity="batch"; call it once per epoch otherwise).

Source code in src/torchtsetlin/models/base.py

def resample_dropout(self) -> None:
    """Draw new clause / literal dropout masks (called automatically per update when
    ``drop_granularity="batch"``; call it once per epoch otherwise)."""
    dev = self.ta_state.device
    if self.drop_clause_p > 0.0:
        self._clause_active = torch.rand(self.n_clauses_total, device=dev) >= self.drop_clause_p
    else:
        self._clause_active = None
    if self.drop_literal_p > 0.0 and self.is_initialized:
        keep = torch.rand(self.n_literals, device=dev) >= self.drop_literal_p
        self._literal_active = keep.to(self.compute_dtype)
    else:
        self._literal_active = None

included_mask

included_mask() -> Tensor

Boolean (C, 2F) mask of included literals (first F columns = positive literals x_k, last F columns = negated literals NOT x_k).

Source code in src/torchtsetlin/models/base.py

def included_mask(self) -> Tensor:
    """Boolean ``(C, 2F)`` mask of included literals (first ``F`` columns = positive
    literals ``x_k``, last ``F`` columns = negated literals ``NOT x_k``)."""
    self._check_initialized()
    return self.ta_state >= self.n_states

clause_literals

clause_literals(clause: int) -> Dict[str, List[int]]

Indices of the included literals of one clause: {"positive": [...], "negated": [...]}.

Source code in src/torchtsetlin/models/base.py

def clause_literals(self, clause: int) -> Dict[str, List[int]]:
    """Indices of the included literals of one clause: ``{"positive": [...], "negated": [...]}``."""
    mask = self.included_mask()[clause]
    Fn = int(self.n_features)  # type: ignore[arg-type]
    pos = torch.nonzero(mask[:Fn], as_tuple=False).flatten().tolist()
    neg = torch.nonzero(mask[Fn:], as_tuple=False).flatten().tolist()
    return {"positive": pos, "negated": neg}

literal_names

literal_names(feature_names: Optional[Sequence[str]] = None) -> List[str]

Names of all 2F literals, e.g. ["x0", "x1", "NOT x0", "NOT x1"].

Source code in src/torchtsetlin/models/base.py

def literal_names(self, feature_names: Optional[Sequence[str]] = None) -> List[str]:
    """Names of all ``2F`` literals, e.g. ``["x0", "x1", "NOT x0", "NOT x1"]``."""
    Fn = int(self.n_features)  # type: ignore[arg-type]
    names = list(feature_names or self.feature_names or self.default_feature_names())
    if len(names) != Fn:
        raise ValueError(f"Expected {Fn} feature names, got {len(names)}")
    return names + [f"NOT {n}" for n in names]

clause_expression

clause_expression(clause: int, feature_names: Optional[Sequence[str]] = None, joiner: str = ' AND ') -> str

Human readable conjunction for one clause ("x1 AND NOT x3"; "TRUE" if empty).

Source code in src/torchtsetlin/models/base.py

def clause_expression(
    self, clause: int, feature_names: Optional[Sequence[str]] = None, joiner: str = " AND "
) -> str:
    """Human readable conjunction for one clause (``"x1 AND NOT x3"``; ``"TRUE"`` if empty)."""
    lits = self.literal_names(feature_names)
    mask = self.included_mask()[clause]
    idx = torch.nonzero(mask, as_tuple=False).flatten().tolist()
    if not idx:
        return "TRUE"
    return joiner.join(lits[i] for i in idx)

literal_frequency

literal_frequency() -> Tensor

How many clauses include each literal -> (2F,) long tensor.

Source code in src/torchtsetlin/models/base.py

def literal_frequency(self) -> Tensor:
    """How many clauses include each literal -> ``(2F,)`` long tensor."""
    return self.included_mask().sum(dim=0)

state_summary

state_summary() -> Dict[str, float]

Quick statistics about the automata (useful for logging).

Source code in src/torchtsetlin/models/base.py

def state_summary(self) -> Dict[str, float]:
    """Quick statistics about the automata (useful for logging)."""
    self._check_initialized()
    counts = self.include_count.to(torch.float32)
    return {
        "included_literals_mean": float(counts.mean()),
        "included_literals_max": float(counts.max()),
        "empty_clauses": int((counts == 0).sum()),
        "n_clauses": int(self.n_clauses_total),
        "n_literals": int(self.n_literals),
    }