torchtsetlin.functional

torchtsetlin.functional

Stateless building blocks of Tsetlin machine learning, expressed as tensor operations.

Everything here works on any device (CPU / CUDA / MPS) and never touches autograd. The model classes in :mod:torchtsetlin.models are thin, stateful wrappers around these functions, so advanced users can compose their own Tsetlin machine variants.

Conventions

• A literal vector is [x, ~x] (length 2F for F Boolean features).
• Each clause owns one Tsetlin automaton (TA) per literal. Its state is an integer in [0, 2N-1]; states >= N mean the literal is included (memorized) in the clause, states < N mean excluded (forgotten). N is n_states.
• Clause output is the AND of its included literals. An empty clause (no included literal) evaluates to 1 during learning and 0 during prediction.
• Feedback probabilities follow the specificity s: memorize probability (s-1)/s, forget probability 1/s (Type Ia/Ib), Type II is deterministic.

to_literals

to_literals(x: Tensor, dtype: dtype = torch.float32) -> Tensor

Turn Boolean features x of shape (..., F) into literals [x, ~x] of shape (..., 2F) with the requested floating dtype (for matmul-based evaluation).

Source code in src/torchtsetlin/functional.py

def to_literals(x: Tensor, dtype: torch.dtype = torch.float32) -> Tensor:
    """Turn Boolean features ``x`` of shape ``(..., F)`` into literals ``[x, ~x]`` of
    shape ``(..., 2F)`` with the requested floating dtype (for matmul-based evaluation)."""
    if x.dtype != torch.bool:
        x = x != 0
    return torch.cat([x, ~x], dim=-1).to(dtype)

clause_violations

clause_violations(literals: Tensor, include: Tensor) -> Tensor

Number of included but False literals per clause.

Parameters:

Name	Type	Description	Default
literals	Tensor	(B, 2F) floating tensor with 0/1 values.	required
include	Tensor	(C, 2F) floating tensor with 0/1 values (1 = literal included).	required

Returns:

Type	Description
Tensor	(B, C) floating tensor; a clause matches an example iff its entry is 0.

Source code in src/torchtsetlin/functional.py

def clause_violations(literals: Tensor, include: Tensor) -> Tensor:
    """Number of *included but False* literals per clause.

    Args:
        literals: ``(B, 2F)`` floating tensor with 0/1 values.
        include: ``(C, 2F)`` floating tensor with 0/1 values (``1`` = literal included).

    Returns:
        ``(B, C)`` floating tensor; a clause matches an example iff its entry is ``0``.
    """
    # (1 - literals) marks False literals; a False literal that is included violates the AND.
    # The literals follow the dtype of ``include`` (the models keep it in float16 on CUDA).
    if literals.dtype != include.dtype:
        literals = literals.to(include.dtype)
    if literals.dim() == 2 and literals.device.type == "cpu" and literals.shape[0] <= 256:
        # clause-major (C, 2F) @ (2F, B) is up to ~2x faster than (B, 2F) @ (2F, C) on the CPU
        # for a few rows (flat mini-batches), but ~3x slower for the thousands of patch rows of
        # a convolutional batch (B=16 MNIST images: 582 vs 187 ms on 16 threads)
        return torch.matmul(include, (1.0 - literals).transpose(0, 1)).transpose(0, 1).contiguous()
    return torch.matmul(1.0 - literals, include.transpose(0, 1))

clause_outputs

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

Evaluate all clauses on a batch of literal vectors.

Parameters:

Name	Type	Description	Default
literals	Tensor	(B, 2F) 0/1 floating tensor.	required
include	Tensor	(C, 2F) 0/1 floating tensor of included literals.	required
include_count	Optional[Tensor]	optional (C,) number of included literals per clause; needed when empty_value is False.	None
empty_value	bool	value of a clause without included literals (True while learning, False when predicting).	True

Returns:

Type	Description
Tensor	(B, C) boolean tensor of clause outputs.

Source code in src/torchtsetlin/functional.py

def clause_outputs(
    literals: Tensor,
    include: Tensor,
    include_count: Optional[Tensor] = None,
    empty_value: bool = True,
) -> Tensor:
    """Evaluate all clauses on a batch of literal vectors.

    Args:
        literals: ``(B, 2F)`` 0/1 floating tensor.
        include: ``(C, 2F)`` 0/1 floating tensor of included literals.
        include_count: optional ``(C,)`` number of included literals per clause; needed
            when ``empty_value`` is ``False``.
        empty_value: value of a clause without included literals (``True`` while
            learning, ``False`` when predicting).

    Returns:
        ``(B, C)`` boolean tensor of clause outputs.
    """
    out = clause_violations(literals, include) == 0
    if not empty_value:
        if include_count is None:
            include_count = include.sum(dim=1)
        out = out & (include_count > 0).unsqueeze(0)
    return out

vote_sums

vote_sums(clause_out: Tensor, weights: Tensor, T: Optional[float] = None) -> Tensor

Weighted vote (class) sums clause_out @ weights optionally clamped to [-T, T].

Parameters:

Name	Type	Description	Default
clause_out	Tensor	(B, C) boolean / 0-1 tensor.	required
weights	Tensor	(C, K) integer or floating weights (K outputs).	required
T	Optional[float]	vote margin; if given the sums are clamped to [-T, T].	None

Source code in src/torchtsetlin/functional.py

def vote_sums(clause_out: Tensor, weights: Tensor, T: Optional[float] = None) -> Tensor:
    """Weighted vote (class) sums ``clause_out @ weights`` optionally clamped to ``[-T, T]``.

    Args:
        clause_out: ``(B, C)`` boolean / 0-1 tensor.
        weights: ``(C, K)`` integer or floating weights (``K`` outputs).
        T: vote margin; if given the sums are clamped to ``[-T, T]``.
    """
    votes = torch.matmul(clause_out.to(torch.float32), weights.to(torch.float32))
    if T is not None:
        votes = votes.clamp_(-float(T), float(T))
    return votes

feedback_probabilities

feedback_probabilities(votes: Tensor, T: float, target: bool) -> Tensor

Probability of giving feedback to a clause given the (clamped) vote sum.

• For the target class the probability is (T - v) / (2T): far below the margin -> update aggressively, at the margin -> stop.
• For a negative class it is (T + v) / (2T).

Source code in src/torchtsetlin/functional.py

def feedback_probabilities(votes: Tensor, T: float, target: bool) -> Tensor:
    """Probability of giving feedback to a clause given the (clamped) vote sum.

    * For the **target** class the probability is ``(T - v) / (2T)``: far below the margin
      -> update aggressively, at the margin -> stop.
    * For a **negative** class it is ``(T + v) / (2T)``.
    """
    v = votes.clamp(-float(T), float(T))
    if target:
        return (float(T) - v) / (2.0 * float(T))
    return (float(T) + v) / (2.0 * float(T))

sample_negative_classes

sample_negative_classes(y: Tensor, n_classes: int, votes: Optional[Tensor] = None, focused: bool = False, generator: Optional[Generator] = None, T: Optional[float] = None) -> Tensor

Sample one negative (non-target) class per example.

With focused=True ("focused negative sampling") a non-target class k is sampled with probability proportional to its negative-feedback probability (T + clip(v_k, -T, T)) / 2T, as in tmu and green-tsetlin, so the classes the machine confuses with the target are corrected first. This needs the vote margin T. Otherwise sampling is uniform.

Source code in src/torchtsetlin/functional.py

def sample_negative_classes(
    y: Tensor,
    n_classes: int,
    votes: Optional[Tensor] = None,
    focused: bool = False,
    generator: Optional[torch.Generator] = None,
    T: Optional[float] = None,
) -> Tensor:
    """Sample one negative (non-target) class per example.

    With ``focused=True`` ("focused negative sampling") a non-target class ``k`` is sampled
    with probability proportional to its negative-feedback probability
    ``(T + clip(v_k, -T, T)) / 2T``, as in tmu and green-tsetlin, so the classes the machine
    confuses with the target are corrected first. This needs the vote margin ``T``.
    Otherwise sampling is uniform.
    """
    if n_classes < 2:
        raise ValueError("Need at least two classes to sample a negative class.")
    if focused and votes is not None:
        if T is None:
            raise ValueError("focused negative sampling needs the vote margin T.")
        w = feedback_probabilities(votes.detach().to(torch.float32), T, target=False)
        w = w + 1e-16  # as green-tsetlin: a class at -T keeps a tiny chance
        w = w.scatter(1, y.view(-1, 1), 0.0)  # never the target itself
        return torch.multinomial(w, 1, generator=generator).squeeze(1)
    offset = torch.randint(1, n_classes, (y.shape[0],), device=y.device, generator=generator)
    return (y + offset) % n_classes

type_i_counts

type_i_counts(literals: Tensor, sel_fire: Tensor, sel_nofire: Tensor) -> Tuple[Tensor, Tensor, Tensor]

Aggregate Type I feedback events over a batch.

Parameters:

Name	Type	Description	Default
literals	Tensor	(B, 2F) 0/1 floating tensor.	required
sel_fire	Tensor	(B, C) 0/1 floating tensor; example b gives Type I feedback to a clause j that matched it (Type Ia).	required
sel_nofire	Tensor	(B, C) 0/1 floating tensor; Type I feedback to a clause that did not match (Type Ib).	required

Returns:

Type	Description
Tensor	(n_true, n_false, n_ib) where n_true[j,k] counts Type Ia events with literal
Tensor	k True (candidates for memorization), n_false[j,k] counts Type Ia events with
Tensor	literal k False (candidates for forgetting) and n_ib[j] counts Type Ib events
Tuple[Tensor, Tensor, Tensor]	(all literals are candidates for forgetting).

Source code in src/torchtsetlin/functional.py

def type_i_counts(
    literals: Tensor, sel_fire: Tensor, sel_nofire: Tensor
) -> Tuple[Tensor, Tensor, Tensor]:
    """Aggregate Type I feedback events over a batch.

    Args:
        literals: ``(B, 2F)`` 0/1 floating tensor.
        sel_fire: ``(B, C)`` 0/1 floating tensor; example ``b`` gives Type I feedback to a
            clause ``j`` **that matched** it (Type Ia).
        sel_nofire: ``(B, C)`` 0/1 floating tensor; Type I feedback to a clause that did
            **not** match (Type Ib).

    Returns:
        ``(n_true, n_false, n_ib)`` where ``n_true[j,k]`` counts Type Ia events with literal
        ``k`` True (candidates for memorization), ``n_false[j,k]`` counts Type Ia events with
        literal ``k`` False (candidates for forgetting) and ``n_ib[j]`` counts Type Ib events
        (all literals are candidates for forgetting).
    """
    sel_fire_t = sel_fire.transpose(0, 1)  # (C, B)
    n_true = torch.matmul(sel_fire_t, literals)  # (C, 2F)
    n_false = torch.matmul(sel_fire_t, 1.0 - literals)
    n_ib = sel_nofire.sum(dim=0)  # (C,)
    return n_true, n_false, n_ib

type_ii_counts

type_ii_counts(literals: Tensor, sel_fire: Tensor) -> Tensor

Aggregate Type II feedback events: n2[j,k] counts (example, clause) pairs where the clause matched the example, received Type II feedback and literal k was False.

Source code in src/torchtsetlin/functional.py

def type_ii_counts(literals: Tensor, sel_fire: Tensor) -> Tensor:
    """Aggregate Type II feedback events: ``n2[j,k]`` counts (example, clause) pairs where the
    clause matched the example, received Type II feedback and literal ``k`` was False."""
    return torch.matmul(sel_fire.transpose(0, 1), 1.0 - literals)

set_compile

set_compile(enabled: bool) -> None

Enable or disable the torch.compile kernels of the update (default: enabled, or the environment variable TORCHTSETLIN_COMPILE=0 to disable). The compiled and eager paths draw from the same distributions but consume random numbers differently, so a seeded run reproduces only under the same setting.

Source code in src/torchtsetlin/functional.py

def set_compile(enabled: bool) -> None:
    """Enable or disable the ``torch.compile`` kernels of the update (default: enabled, or
    the environment variable ``TORCHTSETLIN_COMPILE=0`` to disable). The compiled and eager
    paths draw from the same distributions but consume random numbers differently, so a seeded
    run reproduces only under the same setting."""
    _COMPILE["enabled"] = bool(enabled)

apply_feedback

apply_feedback(state: Tensor, n_true: Tensor, n_false: Tensor, n_ib: Tensor, n2: Tensor, *, n_states: int, s: float, boost_true_positive: bool = True, max_included_literals: Optional[int] = None, include_count: Optional[Tensor] = None, literal_active: Optional[Tensor] = None, generator: Optional[Generator] = None, budget_mode: str = 'soft') -> Tensor

Turn aggregated feedback counts into Tsetlin automata transitions (in place).

Parameters:

Name	Type	Description	Default
state	Tensor	(C, 2F) integer TA states, modified in place.	required
n_true	Tensor	(C, 2F) Type Ia events with the literal True (from :func:type_i_counts).	required
n_false	Tensor	(C, 2F) Type Ia events with the literal False.	required
n_ib	Tensor	(C,) Type Ib events per clause.	required
n2	Tensor	(C, 2F) Type II events with the literal False (from :func:type_ii_counts). All four count tensors are used as scratch space (modified in place) — do not reuse them afterwards.	required
n_states	int	N states per action (include iff state >= N).	required
s	float	specificity; memorize probability (s-1)/s and forget probability 1/s.	required
boost_true_positive	bool	if True the memorization of True literals in matching clauses is deterministic (probability 1) instead of (s-1)/s.	True
max_included_literals	Optional[int]	clause size budget b. A matching clause that includes more than b literals gets no Type Ia feedback; all its Type I events act as Type Ib (forgetting) until it shrinks back (Abeyrathna et al., 2023). Under budget_mode="hard" the budget is also a hard limit on growth.	None
include_count	Optional[Tensor]	(C,) included literals per clause (needed for the constraint).	None
literal_active	Optional[Tensor]	optional (2F,) 0/1 mask; inactive literals receive no feedback (literal dropout).	None
generator	Optional[Generator]	optional RNG.	None
budget_mode	str	"soft" (default) is the reference implementation's rule (tmu): the budget only gates Type Ia, so a clause may overshoot b and is then pruned by Type Ib until it shrinks back. "hard" also reverts new inclusions beyond the budget, so a clause never ends a call above b, however many events were aggregated into it.	'soft'

Returns:

Type	Description
Tensor	The updated state tensor (same object).

Source code in src/torchtsetlin/functional.py

def apply_feedback(
    state: Tensor,
    n_true: Tensor,
    n_false: Tensor,
    n_ib: Tensor,
    n2: Tensor,
    *,
    n_states: int,
    s: float,
    boost_true_positive: bool = True,
    max_included_literals: Optional[int] = None,
    include_count: Optional[Tensor] = None,
    literal_active: Optional[Tensor] = None,
    generator: Optional[torch.Generator] = None,
    budget_mode: str = "soft",
) -> Tensor:
    """Turn aggregated feedback counts into Tsetlin automata transitions (in place).

    Args:
        state: ``(C, 2F)`` integer TA states, modified in place.
        n_true: ``(C, 2F)`` Type Ia events with the literal True (from :func:`type_i_counts`).
        n_false: ``(C, 2F)`` Type Ia events with the literal False.
        n_ib: ``(C,)`` Type Ib events per clause.
        n2: ``(C, 2F)`` Type II events with the literal False (from :func:`type_ii_counts`).
            All four count tensors **are used as scratch space** (modified in place) — do not
            reuse them afterwards.
        n_states: ``N`` states per action (include iff ``state >= N``).
        s: specificity; memorize probability ``(s-1)/s`` and forget probability ``1/s``.
        boost_true_positive: if ``True`` the memorization of True literals in matching
            clauses is deterministic (probability 1) instead of ``(s-1)/s``.
        max_included_literals: clause size budget ``b``. A matching clause that includes more
            than ``b`` literals gets no Type Ia feedback; all its Type I events act as Type Ib
            (forgetting) until it shrinks back (Abeyrathna et al., 2023). Under
            ``budget_mode="hard"`` the budget is also a hard limit on growth.
        include_count: ``(C,)`` included literals per clause (needed for the constraint).
        literal_active: optional ``(2F,)`` 0/1 mask; inactive literals receive no feedback
            (literal dropout).
        generator: optional RNG.
        budget_mode: ``"soft"`` (default) is the reference implementation's rule (tmu): the
            budget only gates Type Ia, so a clause may overshoot ``b`` and is then pruned by
            Type Ib until it shrinks back. ``"hard"`` also reverts new inclusions beyond the
            budget, so a clause never ends a call above ``b``, however many events were
            aggregated into it.

    Returns:
        The updated ``state`` tensor (same object).
    """
    N = int(n_states)
    p_forget = 1.0 / float(s)
    p_mem = 1.0 if boost_true_positive else (float(s) - 1.0) / float(s)

    if max_included_literals is not None:
        if include_count is None:
            include_count = (state >= N).sum(dim=1)
        oversize = (include_count > int(max_included_literals)).unsqueeze(1)
        if bool(oversize.any()):
            # Every Type Ia event of an oversized clause becomes a Type Ib event.
            n_ia_events = n_true[:, :1] + n_false[:, :1]  # identical for every literal
            n_ib = n_ib + (n_ia_events * oversize).squeeze(1)
            keep = (~oversize).to(n_true.dtype)
            n_true.mul_(keep)
            n_false.mul_(keep)

    # Type Ia – memorize True literals of matching clauses.
    inc = _binomial(n_true, p_mem, generator)

    # Type Ia (False literals of matching clauses) + Type Ib (all literals of non-matching
    # clauses) – forget with probability 1/s. Both are Bernoulli(1/s) events, so the sum of
    # counts is again binomial.
    n_false.add_(n_ib.unsqueeze(1))
    dec = _binomial(n_false, p_forget, generator)

    # Type II – deterministically push excluded False literals towards inclusion. Each event
    # moves the automaton one step, but never beyond the include boundary (as in the
    # sequential algorithm, where the clause stops matching once the literal is included).
    room = torch.sub(N, state.to(torch.int32)).clamp_(min=0)
    inc2 = torch.minimum(n2, room.to(n2.dtype), out=n2)

    delta = inc.sub_(dec).add_(inc2)
    if literal_active is not None:
        delta.mul_(literal_active.to(delta.dtype).unsqueeze(0))
    # Add in int32 and clamp before casting back, so that narrow state dtypes (uint8, int16)
    # never see a negative or overflowing intermediate.
    # The new state is formed in the (floating) delta buffer and clamped before it is written
    # back, so narrow state dtypes (uint8, int16) never see a negative or overflowing value.
    # All values are small integers, exactly representable in the compute dtype.
    hard = max_included_literals is not None and budget_mode != "soft"
    was_included = state >= N if hard else None
    new = delta.add_(state).clamp_(0, 2 * N - 1)
    if hard:
        _cap_new_inclusions(new, was_included, N, int(max_included_literals), generator)
    state.copy_(new)
    return state

sample_forget_events

sample_forget_events(counts: Tensor, n_literals: int, p: float, generator: Optional[Generator] = None) -> Tuple[Tensor, Tensor]

Sparse exact draw of per-literal forgetting for row-broadcast event counts.

Row c holds counts[c] Bernoulli(p) trials for each of its n_literals literals (Type Ib feedback applies the same count to every literal of a clause). Every literal therefore needs a Binomial(counts[c], p) decrement. Instead of one draw per literal, the successes of each row's counts[c] * n_literals trials are generated by geometric gaps, so the work scales with the number of successes. Trial t of row c belongs to literal t // counts[c].

Parameters:

Name	Type	Description	Default
counts	Tensor	(C,) non-negative integer-valued counts (any numeric dtype).	required
n_literals	int	literals per row L.	required
p	float	success probability, 0 < p <= 0.5 (the gaps are long only for small p).	required
generator	Optional[Generator]	optional RNG.	None

Returns:

Type	Description
Tensor	(pos, cnt): sorted unique flat positions c * L + l with at least one success
Tensor	and the number of successes at each (int64).

Source code in src/torchtsetlin/functional.py

def sample_forget_events(
    counts: Tensor, n_literals: int, p: float, generator: Optional[torch.Generator] = None
) -> Tuple[Tensor, Tensor]:
    """Sparse exact draw of per-literal forgetting for row-broadcast event counts.

    Row ``c`` holds ``counts[c]`` Bernoulli(``p``) trials for each of its ``n_literals``
    literals (Type Ib feedback applies the same count to every literal of a clause). Every
    literal therefore needs a ``Binomial(counts[c], p)`` decrement. Instead of one draw per
    literal, the successes of each row's ``counts[c] * n_literals`` trials are generated by
    geometric gaps, so the work scales with the number of successes. Trial ``t`` of row ``c``
    belongs to literal ``t // counts[c]``.

    Args:
        counts: ``(C,)`` non-negative integer-valued counts (any numeric dtype).
        n_literals: literals per row ``L``.
        p: success probability, ``0 < p <= 0.5`` (the gaps are long only for small ``p``).
        generator: optional RNG.

    Returns:
        ``(pos, cnt)``: sorted unique flat positions ``c * L + l`` with at least one success
        and the number of successes at each (``int64``).
    """
    dev = counts.device
    L = int(n_literals)
    rows = (counts > 0).nonzero().squeeze(1)
    empty = torch.zeros(0, dtype=torch.long, device=dev)
    if rows.numel() == 0 or p <= 0.0:
        return empty, empty
    n = counts[rows].round().long()
    trials = n * L
    log_q = math.log1p(-float(p))
    mu = trials.double() * float(p)
    # Enough gaps per row that running out is rare; rows that do run out continue below.
    k = (mu + 6.0 * mu.sqrt() + 8.0).ceil().long()
    starts = torch.zeros_like(trials)  # first unexamined trial of each row
    pos_parts, todo = [], torch.arange(rows.numel(), device=dev)
    while todo.numel():
        kk = k[todo]
        row_of = torch.repeat_interleave(torch.arange(todo.numel(), device=dev), kk)
        u = torch.rand(row_of.numel(), device=dev, generator=generator, dtype=torch.float64)
        step = (torch.log(u) / log_q).floor_().long().add_(1)  # trials up to the next success
        csum = step.cumsum(0)
        seg_end = kk.cumsum(0) - 1
        seg_base = csum[seg_end] - kk.new_zeros(kk.shape).index_add_(0, row_of, step)
        t = csum - seg_base[row_of] - 1 + starts[todo][row_of]
        lim = trials[todo][row_of]
        ok = t < lim
        pos_parts.append(rows[todo][row_of][ok] * L + t[ok] // n[todo][row_of][ok])
        last = t[seg_end]
        more = last < trials[todo] - 1
        starts[todo] = last + 1
        todo = todo[more]
        if todo.numel():
            k[todo] = k[todo] * 2
    pos = pos_parts[0] if len(pos_parts) == 1 else torch.sort(torch.cat(pos_parts)).values
    if pos.numel() == 0:
        return empty, empty
    return torch.unique_consecutive(pos, return_counts=True)

apply_feedback_rows

apply_feedback_rows(state: Tensor, n_true: Tensor, n_false: Tensor, n_ib: Tensor, n2: Tensor, rows: Tensor, *, n_states: int, s: float, boost_true_positive: bool = True, max_included_literals: Optional[int] = None, include_count: Optional[Tensor] = None, literal_active: Optional[Tensor] = None, generator: Optional[Generator] = None, budget_mode: str = 'soft') -> Optional[Tuple[Tensor, Tensor]]

:func:apply_feedback for count tensors that are zero outside rows.

Type Ia and Type II counts only exist for clauses that matched and were selected, a small fraction of all clauses. Those rows are updated as a dense slab; the Type Ib forgetting of every other clause is drawn sparsely with :func:sample_forget_events. The transitions follow the same distribution as :func:apply_feedback. The count tensors are read, not modified.

Parameters:

Name	Type	Description	Default
rows	Tensor	(C,) bool mask of the clauses with Type Ia or Type II counts. The other arguments are as in :func:apply_feedback.	required

Returns:

Type	Description
Optional[Tuple[Tensor, Tensor]]	(slab_rows, other_pos): the clause indices updated as a dense slab and the flat
Optional[Tuple[Tensor, Tensor]]	positions changed elsewhere, so that cached include masks can be refreshed
Optional[Tuple[Tensor, Tensor]]	incrementally. None when the dense path was used (refresh everything).

Source code in src/torchtsetlin/functional.py

def apply_feedback_rows(
    state: Tensor,
    n_true: Tensor,
    n_false: Tensor,
    n_ib: Tensor,
    n2: Tensor,
    rows: Tensor,
    *,
    n_states: int,
    s: float,
    boost_true_positive: bool = True,
    max_included_literals: Optional[int] = None,
    include_count: Optional[Tensor] = None,
    literal_active: Optional[Tensor] = None,
    generator: Optional[torch.Generator] = None,
    budget_mode: str = "soft",
) -> Optional[Tuple[Tensor, Tensor]]:
    """:func:`apply_feedback` for count tensors that are zero outside ``rows``.

    Type Ia and Type II counts only exist for clauses that matched and were selected, a small
    fraction of all clauses. Those rows are updated as a dense slab; the Type Ib forgetting of
    every other clause is drawn sparsely with :func:`sample_forget_events`. The transitions
    follow the same distribution as :func:`apply_feedback`. The count tensors are read, not
    modified.

    Args:
        rows: ``(C,)`` bool mask of the clauses with Type Ia or Type II counts. The other
            arguments are as in :func:`apply_feedback`.

    Returns:
        ``(slab_rows, other_pos)``: the clause indices updated as a dense slab and the flat
        positions changed elsewhere, so that cached include masks can be refreshed
        incrementally. ``None`` when the dense path was used (refresh everything).
    """
    N = int(n_states)
    p_forget = 1.0 / float(s)
    if p_forget > 0.5:  # forgetting is not sparse; the dense update is as fast
        apply_feedback(state, n_true.clone(), n_false.clone(), n_ib.clone(), n2.clone(),
                       n_states=N, s=s, boost_true_positive=boost_true_positive,
                       max_included_literals=max_included_literals, include_count=include_count,
                       literal_active=literal_active, generator=generator,
                       budget_mode=budget_mode)
        return None
    p_mem = 1.0 if boost_true_positive else (float(s) - 1.0) / float(s)
    C, L = state.shape
    idx = rows.nonzero().squeeze(1)
    nt, nf, n2s = n_true[idx], n_false[idx], n2[idx]
    nib = n_ib.clone()
    if max_included_literals is not None and idx.numel():
        if include_count is None:
            include_count = (state >= N).sum(dim=1)
        over = include_count[idx] > int(max_included_literals)
        if bool(over.any()):
            # Every Type Ia event of an oversized clause becomes a Type Ib event.
            nib.index_add_(0, idx[over], (nt[over, 0] + nf[over, 0]).to(nib.dtype))
            keep = (~over).to(nt.dtype).unsqueeze(1)
            nt.mul_(keep)
            nf.mul_(keep)

    # Type Ib forgetting of the clauses outside the slab: one compiled dense pass over their
    # rows on the CPU, else an exact sparse draw of the forget events.
    other = None
    if state.device.type == "cpu" and _COMPILE["enabled"] and generator is None:
        ib_rows = ((nib > 0) & ~rows).nonzero().squeeze(1)
        if _compile_wanted(ib_rows.numel() * L):
            other = _forget_rows(state, ib_rows, nib[ib_rows], p_forget, literal_active, N)
    if other is None:
        pos, cnt = sample_forget_events(nib, L, p_forget, generator)
        if literal_active is not None and pos.numel():
            alive = literal_active.reshape(-1)[pos % L] > 0
            pos, cnt = pos[alive], cnt[alive]
        slab_of = torch.full((C,), -1, dtype=torch.long, device=state.device)
        slab_of[idx] = torch.arange(idx.numel(), device=state.device)
        s_i = slab_of[pos // L]
        in_slab = s_i >= 0

    if idx.numel():
        inc = _binomial(nt, p_mem, generator)
        if other is not None:  # the slab's Type Ib events are part of its dense draw
            dec = _binomial(nf + nib[idx].unsqueeze(1), p_forget, generator)
        else:
            dec = _binomial(nf, p_forget, generator)
            if bool(in_slab.any()):
                dec.view(-1).index_add_(0, s_i[in_slab] * L + pos[in_slab] % L,
                                        cnt[in_slab].to(dec.dtype))
        sub = state[idx].to(torch.int32)
        room = torch.sub(N, sub).clamp_(min=0)
        inc2 = torch.minimum(n2s, room.to(n2s.dtype))
        delta = inc.sub_(dec).add_(inc2)
        if literal_active is not None:
            delta.mul_(literal_active.to(delta.dtype).unsqueeze(0))
        was_included = sub >= N
        sub.add_(delta.to(torch.int32)).clamp_(0, 2 * N - 1)
        if max_included_literals is not None and budget_mode != "soft":
            _cap_new_inclusions(sub, was_included, N, int(max_included_literals), generator)
        state[idx] = sub.to(state.dtype)

    if other is not None:
        return idx, other
    other = pos[~in_slab]
    if other.numel():
        flat = state.view(-1)
        vals = flat[other].to(torch.int32).sub_(cnt[~in_slab].to(torch.int32)).clamp_(min=0)
        flat[other] = vals.to(state.dtype)
    return idx, other

commit_single_example

commit_single_example(state: Tensor, include: Tensor, include_count: Tensor, literals: Tensor, fire_i: Tensor, nofire_i: Tensor, fire_ii: Tensor, *, n_states: int, s: float, boost_true_positive: bool = True, max_included_literals: Optional[int] = None, literal_active: Optional[Tensor] = None, literals_ii: Optional[Tensor] = None) -> Optional[Tensor]

The feedback of one example, applied by a single fused torch.compile kernel.

Equivalent in distribution to accumulating the example's counts (the outer products of fire_i / nofire_i / fire_ii with the literals) and calling :func:apply_feedback with the soft budget, but without materialising any (C, 2F) count tensor and without host synchronisation. Used by sequential mode on CUDA.

Parameters:

Name	Type	Description	Default
state	Tensor	(C, 2F) TA states, updated in place.	required
include	Tensor	(C, 2F) cached include mask, updated in place.	required
include_count	Tensor	(C,) included literals per clause before the update; updated in place.	required
literals	Tensor	(2F,) 0/1 literals of the example, or (C, 2F) literals per clause (a convolutional model updates each clause from its own random matching patch).	required
fire_i	Tensor	(C,) 0/1 mask of the clauses that get Type Ia feedback.	required
nofire_i	Tensor	(C,) 0/1 mask of the clauses that get Type Ib feedback.	required
fire_ii	Tensor	(C,) 0/1 mask of the clauses that get Type II feedback.	required
n_states	int	as in :func:apply_feedback.	required
s	float	as in :func:apply_feedback.	required
boost_true_positive	bool	as in :func:apply_feedback.	True
max_included_literals	Optional[int]	as in :func:apply_feedback (soft budget).	None
literal_active	Optional[Tensor]	as in :func:apply_feedback.	None
literals_ii	Optional[Tensor]	literals of the Type II events, shaped like literals; defaults to literals (a convolutional model draws the patch of each kind independently).	None

Returns:

Type	Description
Optional[Tensor]	The new (C,) include counts, or None when the compiled kernel is unavailable
Optional[Tensor]	(nothing has been modified; the caller falls back to :func:apply_feedback_rows).

Source code in src/torchtsetlin/functional.py

def commit_single_example(
    state: Tensor,
    include: Tensor,
    include_count: Tensor,
    literals: Tensor,
    fire_i: Tensor,
    nofire_i: Tensor,
    fire_ii: Tensor,
    *,
    n_states: int,
    s: float,
    boost_true_positive: bool = True,
    max_included_literals: Optional[int] = None,
    literal_active: Optional[Tensor] = None,
    literals_ii: Optional[Tensor] = None,
) -> Optional[Tensor]:
    """The feedback of **one** example, applied by a single fused ``torch.compile`` kernel.

    Equivalent in distribution to accumulating the example's counts (the outer products of
    ``fire_i`` / ``nofire_i`` / ``fire_ii`` with the literals) and calling
    :func:`apply_feedback` with the soft budget, but without materialising any ``(C, 2F)``
    count tensor and without host synchronisation. Used by sequential mode on CUDA.

    Args:
        state: ``(C, 2F)`` TA states, updated in place.
        include: ``(C, 2F)`` cached include mask, updated in place.
        include_count: ``(C,)`` included literals per clause *before* the update; updated in
            place.
        literals: ``(2F,)`` 0/1 literals of the example, or ``(C, 2F)`` literals per clause
            (a convolutional model updates each clause from its own random matching patch).
        fire_i: ``(C,)`` 0/1 mask of the clauses that get Type Ia feedback.
        nofire_i: ``(C,)`` 0/1 mask of the clauses that get Type Ib feedback.
        fire_ii: ``(C,)`` 0/1 mask of the clauses that get Type II feedback.
        n_states: as in :func:`apply_feedback`.
        s: as in :func:`apply_feedback`.
        boost_true_positive: as in :func:`apply_feedback`.
        max_included_literals: as in :func:`apply_feedback` (soft budget).
        literal_active: as in :func:`apply_feedback`.
        literals_ii: literals of the Type II events, shaped like ``literals``; defaults to
            ``literals`` (a convolutional model draws the patch of each kind independently).

    Returns:
        The new ``(C,)`` include counts, or ``None`` when the compiled kernel is unavailable
        (nothing has been modified; the caller falls back to :func:`apply_feedback_rows`).
    """
    lit = literals if literals.dim() == 2 else literals.unsqueeze(0)
    lit_ii = lit if literals_ii is None else (
        literals_ii if literals_ii.dim() == 2 else literals_ii.unsqueeze(0))
    N = int(n_states)
    p_mem = 1.0 if boost_true_positive else (float(s) - 1.0) / float(s)
    budget = -1 if max_included_literals is None else int(max_included_literals)
    active = None if literal_active is None else literal_active.reshape(1, -1).to(torch.float32)
    return _run_compiled(_single_example_kernel, state, include, lit, lit_ii, fire_i, nofire_i,
                         fire_ii, include_count, active, N, p_mem, 1.0 / float(s), budget)

commit_dense

commit_dense(state: Tensor, include: Tensor, include_count: Tensor, n_true: Tensor, n_false: Tensor, n_ib: Tensor, n2: Tensor, *, n_states: int, s: float, boost_true_positive: bool = True, max_included_literals: Optional[int] = None, literal_active: Optional[Tensor] = None, max_events: Optional[int] = None) -> Optional[Tensor]

:func:apply_feedback (soft budget) as one fused torch.compile kernel that also refreshes the include cache.

The binomials are drawn inside the kernel by inverse-CDF lookup, so the update reads the counts and the state once and writes the state, the include mask and zeroed counts once, instead of the ~20 full passes of the eager update. Used by batched updates on CUDA.

Parameters:

Name	Type	Description	Default
state	Tensor	(C, 2F) TA states, updated in place.	required
include	Tensor	(C, 2F) cached include mask, updated in place.	required
include_count	Tensor	(C,) included literals per clause before the update; updated in place.	required
n_true	Tensor	as in :func:apply_feedback; zeroed on return (as are n_false and n2; n_ib is left as it was).	required
n_false	Tensor	as in :func:apply_feedback.	required
n_ib	Tensor	as in :func:apply_feedback.	required
n2	Tensor	as in :func:apply_feedback.	required
n_states	int	as in :func:apply_feedback.	required
s	float	as in :func:apply_feedback.	required
boost_true_positive	bool	as in :func:apply_feedback.	True
max_included_literals	Optional[int]	as in :func:apply_feedback (soft budget).	None
literal_active	Optional[Tensor]	as in :func:apply_feedback.	None
max_events	Optional[int]	an upper bound on any count (e.g. the examples in the update). With it the table width is chosen without reading the counts back to the host.	None

Returns:

Type	Description
Optional[Tensor]	The new (C,) include counts (include_count itself), or None (nothing
Optional[Tensor]	modified) when the kernel cannot be used: compilation disabled or failed, or more than
Optional[Tensor]	_WIDE_TABLE_N - 1 events aggregated into one automaton. The caller then takes the
Optional[Tensor]	eager path.

Source code in src/torchtsetlin/functional.py

def commit_dense(
    state: Tensor,
    include: Tensor,
    include_count: Tensor,
    n_true: Tensor,
    n_false: Tensor,
    n_ib: Tensor,
    n2: Tensor,
    *,
    n_states: int,
    s: float,
    boost_true_positive: bool = True,
    max_included_literals: Optional[int] = None,
    literal_active: Optional[Tensor] = None,
    max_events: Optional[int] = None,
) -> Optional[Tensor]:
    """:func:`apply_feedback` (soft budget) as one fused ``torch.compile`` kernel that also
    refreshes the include cache.

    The binomials are drawn inside the kernel by inverse-CDF lookup, so the update reads the
    counts and the state once and writes the state, the include mask and zeroed counts once,
    instead of the ~20 full passes of the eager update. Used by batched updates on CUDA.

    Args:
        state: ``(C, 2F)`` TA states, updated in place.
        include: ``(C, 2F)`` cached include mask, updated in place.
        include_count: ``(C,)`` included literals per clause *before* the update; updated in
            place.
        n_true: as in :func:`apply_feedback`; zeroed on return (as are ``n_false`` and
            ``n2``; ``n_ib`` is left as it was).
        n_false: as in :func:`apply_feedback`.
        n_ib: as in :func:`apply_feedback`.
        n2: as in :func:`apply_feedback`.
        n_states: as in :func:`apply_feedback`.
        s: as in :func:`apply_feedback`.
        boost_true_positive: as in :func:`apply_feedback`.
        max_included_literals: as in :func:`apply_feedback` (soft budget).
        literal_active: as in :func:`apply_feedback`.
        max_events: an upper bound on any count (e.g. the examples in the update). With it the
            table width is chosen without reading the counts back to the host.

    Returns:
        The new ``(C,)`` include counts (``include_count`` itself), or ``None`` (nothing
        modified) when the kernel cannot be used: compilation disabled or failed, or more than
        ``_WIDE_TABLE_N - 1`` events aggregated into one automaton. The caller then takes the
        eager path.
    """
    p_mem = 1.0 if boost_true_positive else (float(s) - 1.0) / float(s)
    if (not _compile_wanted(state.numel()) or _dense_commit_kernel in _FAILED
            or p_mem <= 0.0):
        return None
    # Type Ia events per clause are n_true + n_false of any one literal; with Type Ib they
    # bound every count that is sampled.
    if max_events is not None and int(max_events) < _WIDE_TABLE_N:
        most = int(max_events)
    else:  # no bound given (or a large one): read the actual largest count back
        events = n_true[:, 0] + n_false[:, 0] + n_ib
        most = int(events.max())
    if most >= _WIDE_TABLE_N:
        return None
    width = _bisect_width(most)
    N = int(n_states)
    dev = state.device

    def table(p: float) -> Tuple[Optional[Tensor], bool]:
        if p >= 1.0:
            return None, False
        q = min(p, 1.0 - p)
        return _wide_thresholds(q, dev), p > 0.5

    thr_mem, mem_flip = table(p_mem)
    thr_forget, forget_flip = table(1.0 / float(s))
    budget = -1 if max_included_literals is None else int(max_included_literals)
    active = None if literal_active is None else literal_active.reshape(1, -1).to(torch.float32)
    done = _run_compiled(_dense_commit_kernel, state, include, include_count, n_true, n_false,
                         n_ib, n2, thr_mem, thr_forget, active, N, mem_flip, forget_flip, budget,
                         width)
    if done is None and _dense_commit_kernel in _FAILED:
        return None
    torch.sum(state >= N, dim=1, out=include_count)
    return include_count

predict_proba_from_votes

predict_proba_from_votes(votes: Tensor, T: float, temperature: float = 1.0) -> Tensor

Softmax over vote sums normalised by T – a convenient (heuristic) probability.

votes / T lies in [-1, 1]; dividing by temperature sharpens (<1) or flattens (>1) the distribution.

Source code in src/torchtsetlin/functional.py

def predict_proba_from_votes(votes: Tensor, T: float, temperature: float = 1.0) -> Tensor:
    """Softmax over vote sums normalised by ``T`` – a convenient (heuristic) probability.

    ``votes / T`` lies in ``[-1, 1]``; dividing by ``temperature`` sharpens (<1) or flattens
    (>1) the distribution.
    """
    z = votes.to(torch.float32) / (float(T) * float(temperature))
    return torch.softmax(z, dim=-1)

confidence_from_votes

confidence_from_votes(votes: Tensor, T: float) -> Tensor

Confidence in the predicted class, as in Chapter 7 of An Introduction to Tsetlin Machines: the winning vote sum normalised to [0, 1] by the vote margin, (v_max + T) / (2T). Returns shape (B,).

Source code in src/torchtsetlin/functional.py

def confidence_from_votes(votes: Tensor, T: float) -> Tensor:
    """Confidence in the predicted class, as in Chapter 7 of *An Introduction to Tsetlin
    Machines*: the winning vote sum normalised to ``[0, 1]`` by the vote margin,
    ``(v_max + T) / (2T)``. Returns shape ``(B,)``."""
    v = votes.to(torch.float32).clamp(-float(T), float(T))
    return (v.max(dim=-1).values + float(T)) / (2.0 * float(T))

boolean_pool

boolean_pool(x: Tensor, factor: int = 2, mode: str = 'or') -> Tensor

Pool Boolean planes (B, Z, H, W) by an integer factor.

mode="or" keeps a bit that is set anywhere in the window (the Boolean analogue of max-pooling, and what CTM-UNet uses between blocks); mode="and" keeps only bits set everywhere in it. Both stay Boolean, so the result can be fed straight back into another Tsetlin machine — no float parameter and no gradient anywhere, unlike a strided convolution.

Parameters:

Name	Type	Description	Default
x	Tensor	(B, Z, H, W) Boolean-like tensor.	required
factor	int	pooling window and stride.	2
mode	str	"or" or "and".	'or'

Returns:

Type	Description
Tensor	(B, Z, H // factor, W // factor) bool.

Source code in src/torchtsetlin/functional.py

def boolean_pool(x: Tensor, factor: int = 2, mode: str = "or") -> Tensor:
    """Pool Boolean planes ``(B, Z, H, W)`` by an integer ``factor``.

    ``mode="or"`` keeps a bit that is set anywhere in the window (the Boolean analogue of
    max-pooling, and what *CTM-UNet* uses between blocks); ``mode="and"`` keeps only bits set
    everywhere in it. Both stay Boolean, so the result can be fed straight back into another
    Tsetlin machine — no float parameter and no gradient anywhere, unlike a strided
    convolution.

    Args:
        x: ``(B, Z, H, W)`` Boolean-like tensor.
        factor: pooling window and stride.
        mode: ``"or"`` or ``"and"``.

    Returns:
        ``(B, Z, H // factor, W // factor)`` bool.
    """
    if x.dim() != 4:
        raise ValueError("boolean_pool expects (B, Z, H, W)")
    k = int(factor)
    if k < 1:
        raise ValueError("factor must be >= 1")
    if k == 1:
        return x.to(torch.bool)
    f = x.to(torch.float32)
    if mode == "or":
        return torch.nn.functional.max_pool2d(f, k) > 0.5
    if mode == "and":
        return torch.nn.functional.max_pool2d(-f, k).neg() > 0.5
    raise ValueError("mode must be 'or' or 'and'")

thermometer_cast

thermometer_cast(votes: Tensor, levels: Sequence[float]) -> Tensor

Turn vote sums into Boolean planes without losing their ordering.

A Tsetlin machine can only read Boolean features, so passing one machine's votes to another needs a Booleanization that preserves how strong the vote was. A thermometer does: the planes of v are [v >= l] for each level l, so a larger vote sets a superset of the bits a smaller one does, and a clause can refer to "at least this confident" with one literal.

Parameters:

Name	Type	Description	Default
votes	Tensor	(B, K, H, W) vote map (or (B, K) vote sums).	required
levels	Sequence[float]	ascending thresholds, e.g. (-40, -20, -8, 0, 8, 20, 40).	required

Returns:

Type	Description
Tensor	(B, K * len(levels), H, W) bool for a vote map, (B, K * len(levels)) for
Tensor	vote sums. Channel order is class-major: all levels of class 0, then class 1, ...

Source code in src/torchtsetlin/functional.py

def thermometer_cast(votes: Tensor, levels: Sequence[float]) -> Tensor:
    """Turn vote sums into Boolean planes without losing their ordering.

    A Tsetlin machine can only read Boolean features, so passing one machine's votes to
    another needs a Booleanization that preserves *how strong* the vote was. A thermometer
    does: the planes of ``v`` are ``[v >= l]`` for each level ``l``, so a larger vote sets a
    superset of the bits a smaller one does, and a clause can refer to "at least this
    confident" with one literal.

    Args:
        votes: ``(B, K, H, W)`` vote map (or ``(B, K)`` vote sums).
        levels: ascending thresholds, e.g. ``(-40, -20, -8, 0, 8, 20, 40)``.

    Returns:
        ``(B, K * len(levels), H, W)`` bool for a vote map, ``(B, K * len(levels))`` for
        vote sums. Channel order is class-major: all levels of class 0, then class 1, ...
    """
    lv = torch.as_tensor(list(levels), dtype=votes.dtype, device=votes.device)
    if votes.dim() == 4:
        planes = votes.unsqueeze(2) >= lv.view(1, 1, -1, 1, 1)
        return planes.reshape(votes.shape[0], -1, votes.shape[2], votes.shape[3])
    if votes.dim() == 2:
        return (votes.unsqueeze(2) >= lv.view(1, 1, -1)).reshape(votes.shape[0], -1)
    raise ValueError("thermometer_cast expects (B, K, H, W) or (B, K)")