torchtsetlin.data

Encoders

torchtsetlin.data.encoders

Booleanization: turning continuous, categorical and image data into Boolean features.

All encoders are torch.nn.Module objects so they can be moved to a device, saved in a state_dict and composed with :class:torchtsetlin.data.Compose. They follow the scikit-learn fit / transform protocol on top of forward (transform is an alias of forward; fit_transform chains both).

BooleanEncoder

Bases: Module

Base class for encoders. Subclasses implement :meth:forward (and optionally :meth:fit). transform is an alias of forward.

fit

fit(x, y=None) -> BooleanEncoder

Fit encoder statistics (no-op for stateless encoders).

Source code in src/torchtsetlin/data/encoders.py

def fit(self, x, y=None) -> BooleanEncoder:  # noqa: D401
    """Fit encoder statistics (no-op for stateless encoders)."""
    return self

Binarizer

Binarizer(threshold: Union[float, Tensor] = 0.5, inclusive: bool = False)

Bases: BooleanEncoder

x > threshold (or >=) per element. threshold may be a scalar or a tensor broadcastable to the input (e.g. one threshold per feature).

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, threshold: Union[float, Tensor] = 0.5, inclusive: bool = False) -> None:
    super().__init__()
    self.register_buffer("threshold", torch.as_tensor(threshold, dtype=torch.float32))
    self.inclusive = bool(inclusive)

ThermometerEncoder

ThermometerEncoder(n_bits: int = 8, strategy: str = 'quantile', thresholds: Optional[Tensor] = None, value_range: Optional[Sequence[float]] = None, flatten: bool = True)

Bases: BooleanEncoder

Thermometer (ordinal) encoding of continuous features.

Each feature v becomes n_bits Booleans [v >= t_1, ..., v >= t_n] for increasing thresholds t_k. Clauses can then express ranges with two literals (x>=t_a AND NOT x>=t_b). Thresholds are learned by :meth:fit (quantiles by default) or given explicitly.

Parameters:

Name	Type	Description	Default
n_bits	int	thresholds per feature.	8
strategy	str	threshold strategy (see :func:thermometer_thresholds).	'quantile'
thresholds	Optional[Tensor]	explicit (F, n_bits) thresholds (skips fitting).	None
value_range	Optional[Sequence[float]]	(lo, hi) for the uniform strategy (e.g. (0, 255) for images).	None
flatten	bool	return (N, F*n_bits) instead of (N, F, n_bits).	True

Source code in src/torchtsetlin/data/encoders.py

def __init__(
    self,
    n_bits: int = 8,
    strategy: str = "quantile",
    thresholds: Optional[Tensor] = None,
    value_range: Optional[Sequence[float]] = None,
    flatten: bool = True,
) -> None:
    super().__init__()
    self.n_bits = int(n_bits)
    self.strategy = strategy
    self.value_range = tuple(value_range) if value_range is not None else None
    self.flatten = bool(flatten)
    if thresholds is not None:
        self.register_buffer("thresholds", torch.as_tensor(thresholds, dtype=torch.float32))
    else:
        self.register_buffer("thresholds", torch.empty(0))

feature_names

feature_names(names: Sequence[str]) -> List[str]

Names of the Boolean outputs, e.g. "age>=30.0".

Source code in src/torchtsetlin/data/encoders.py

def feature_names(self, names: Sequence[str]) -> List[str]:
    """Names of the Boolean outputs, e.g. ``"age>=30.0"``."""
    thr = self.thresholds
    if thr.shape[0] == 1:
        thr = thr.expand(len(names), -1)
    return [f"{n}>={float(thr[i, k]):.4g}" for i, n in enumerate(names) for k in range(self.n_bits)]

decode_range

decode_range(feature: int, included_pos: Iterable[int], included_neg: Iterable[int])

Range (lo, hi) implied by included positive/negated bits of one feature.

Source code in src/torchtsetlin/data/encoders.py

def decode_range(self, feature: int, included_pos: Iterable[int], included_neg: Iterable[int]):
    """Range ``(lo, hi)`` implied by included positive/negated bits of one feature."""
    thr = self.thresholds[feature if self.thresholds.shape[0] > 1 else 0]
    lo = max((float(thr[k]) for k in included_pos), default=-float("inf"))
    hi = min((float(thr[k]) for k in included_neg), default=float("inf"))
    return lo, hi

OneHotEncoder

OneHotEncoder(n_categories: Optional[Sequence[int]] = None)

Bases: BooleanEncoder

One Boolean per category value per column. Categories are integers (fit records the number of categories per column, or pass n_categories).

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, n_categories: Optional[Sequence[int]] = None) -> None:
    super().__init__()
    self.register_buffer(
        "n_categories",
        torch.as_tensor(list(n_categories), dtype=torch.long) if n_categories else torch.empty(0, dtype=torch.long),
    )

BitPlaneEncoder

BitPlaneEncoder(n_bits: int = 8, channel_dim: int = 1)

Bases: BooleanEncoder

Plain binary (place-value) encoding of non-negative integers with n_bits bits, most significant bit first. Images (N, C, H, W) become (N, C*n_bits, H, W).

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, n_bits: int = 8, channel_dim: int = 1) -> None:
    super().__init__()
    self.n_bits = int(n_bits)
    self.channel_dim = channel_dim

AdaptiveThresholdEncoder

AdaptiveThresholdEncoder(block_size: int = 11, C: float = 2.0, method: str = 'gaussian')

Bases: BooleanEncoder

Adaptive (local) thresholding for grayscale / multi-channel images, as used to booleanize MNIST/CIFAR for convolutional Tsetlin machines.

bit = pixel > local_mean(pixel) - C where the local mean is computed over a block_size x block_size window, either uniformly (method="mean") or with a Gaussian kernel (method="gaussian", OpenCV ADAPTIVE_THRESH_GAUSSIAN_C).

Input (N, C, H, W) (or (N, H, W)) -> Boolean of the same shape.

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, block_size: int = 11, C: float = 2.0, method: str = "gaussian") -> None:
    super().__init__()
    if block_size % 2 == 0 or block_size < 3:
        raise ValueError("block_size must be odd and >= 3")
    if method not in ("gaussian", "mean"):
        raise ValueError("method must be 'gaussian' or 'mean'")
    self.block_size = int(block_size)
    self.C = float(C)
    self.method = method
    if method == "gaussian":
        sigma = 0.3 * ((block_size - 1) * 0.5 - 1) + 0.8  # OpenCV's default sigma
        ax = torch.arange(block_size, dtype=torch.float32) - (block_size - 1) / 2
        g1 = torch.exp(-(ax**2) / (2 * sigma**2))
        g1 = g1 / g1.sum()
        kernel = torch.outer(g1, g1)
    else:
        kernel = torch.full((block_size, block_size), 1.0 / block_size**2)
    self.register_buffer("kernel", kernel.view(1, 1, block_size, block_size))

ColorThermometerEncoder

ColorThermometerEncoder(n_bits: int = 8, value_range: Sequence[float] = (0.0, 255.0))

Bases: BooleanEncoder

Per-channel thermometer encoding of images: each pixel channel becomes n_bits planes [v >= k * 255/(n_bits+1)]. (N, C, H, W) -> (N, C*n_bits, H, W).

value_range defaults to (0, 255); use (0, 1) for float images.

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, n_bits: int = 8, value_range: Sequence[float] = (0.0, 255.0)) -> None:
    super().__init__()
    lo, hi = float(value_range[0]), float(value_range[1])
    steps = torch.arange(1, n_bits + 1, dtype=torch.float32) / (n_bits + 1)
    self.register_buffer("thresholds", lo + (hi - lo) * steps)
    self.n_bits = int(n_bits)

HypervectorEncoder

HypervectorEncoder(vocab_size: int, dim: int = 1024, n_bits: int = 8, bind_position: bool = False, seed: int = 0)

Bases: BooleanEncoder

Sparse binary hypervector encoding of discrete tokens (Halenka et al., 2024).

Every token id t < vocab_size owns a fixed random hypervector with n_bits ones among dim positions. A sample (a set / sequence of token ids) is bundled with bitwise OR. Optionally each position p is bound by cyclically shifting the token's hypervector by p (bind_position=True), so word order is preserved.

Input: (N, L) long tensor of token ids (padding id -1 is ignored). Output: (N, dim) Boolean.

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, vocab_size: int, dim: int = 1024, n_bits: int = 8, bind_position: bool = False, seed: int = 0) -> None:
    super().__init__()
    g = torch.Generator().manual_seed(seed)
    table = torch.zeros(vocab_size, dim, dtype=torch.bool)
    for t in range(vocab_size):
        idx = torch.randperm(dim, generator=g)[:n_bits]
        table[t, idx] = True
    self.register_buffer("table", table)
    self.dim = int(dim)
    self.bind_position = bool(bind_position)

PyramidEncoder

PyramidEncoder(scales: Sequence[int] = (1, 2, 4), patch: int = 3, mode: str = 'or')

Bases: BooleanEncoder

Stack each pixel's neighbourhood at several resolutions into extra Boolean planes.

A k x k patch is all a dense Tsetlin machine sees of a pixel, and some classes simply are not decidable from it — a patch of building interior and a patch of road look alike once illumination varies; telling them apart needs to know where the horizon is. This is the gap an encoder-decoder closes in a CNN, and it can be closed natively here: OR-pool the Boolean planes to build a pyramid, take the same k x k patch at each level, and resample every level back to full resolution. A clause then reads a pixel's immediate neighbourhood and the coarse structure around it, with no float parameter and no gradient anywhere.

Feed the result to a model with patch_size=1: the neighbourhood is already in the planes, and asking for a second one on top multiplies the feature count for nothing.

Parameters:

Name	Type	Description	Default
scales	Sequence[int]	downsampling factors, ascending. 1 is the full-resolution level.	(1, 2, 4)
patch	int	neighbourhood side length taken at each level.	3
mode	str	pooling mode for the pyramid (see :func:torchtsetlin.functional.boolean_pool).	'or'

Shape

• input: (B, Z, H, W) Boolean.
• output: (B, Z * patch * patch * len(scales), H, W) Boolean.

Example

enc = PyramidEncoder(scales=(1, 2, 4), patch=3) planes = enc(x) # (B, Z93, H, W) model = SegmentationTsetlinMachine(4, 300, T=60, patch_size=1)

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, scales: Sequence[int] = (1, 2, 4), patch: int = 3, mode: str = "or") -> None:
    super().__init__()
    self.scales = tuple(int(s) for s in scales)
    self.patch = int(patch)
    self.mode = str(mode)
    if min(self.scales) < 1 or self.patch < 1:
        raise ValueError("scales and patch must be >= 1")

Flatten

Bases: BooleanEncoder

Flatten everything but the batch dimension.

Compose

Compose(*stages)

Bases: BooleanEncoder

Chain encoders / callables. fit fits each stage on the output of the previous.

Source code in src/torchtsetlin/data/encoders.py

def __init__(self, *stages) -> None:
    super().__init__()
    self.stages = nn.ModuleList([s for s in stages if isinstance(s, nn.Module)])
    self._all = list(stages)

thermometer_thresholds

thermometer_thresholds(x: Tensor, n_bits: int, strategy: str = 'quantile', value_range: Optional[Sequence[float]] = None) -> Tensor

Per-feature thresholds (F, n_bits) for thermometer encoding.

Parameters:

Name	Type	Description	Default
x	Tensor	(N, F) training data (only needed for "quantile"/"unique" or when value_range is None).	required
n_bits	int	thresholds per feature.	required
strategy	str	"quantile" (equal-frequency), "uniform" (equal-width) or "unique" (evenly spaced ranks of the unique training values).	'quantile'
value_range	Optional[Sequence[float]]	optional global (lo, hi) for the uniform strategy.	None

Source code in src/torchtsetlin/data/encoders.py

def thermometer_thresholds(
    x: Tensor,
    n_bits: int,
    strategy: str = "quantile",
    value_range: Optional[Sequence[float]] = None,
) -> Tensor:
    """Per-feature thresholds ``(F, n_bits)`` for thermometer encoding.

    Args:
        x: ``(N, F)`` training data (only needed for ``"quantile"``/``"unique"`` or when
            ``value_range`` is ``None``).
        n_bits: thresholds per feature.
        strategy: ``"quantile"`` (equal-frequency), ``"uniform"`` (equal-width) or
            ``"unique"`` (evenly spaced ranks of the unique training values).
        value_range: optional global ``(lo, hi)`` for the uniform strategy.
    """
    x = x.to(torch.float32)
    Fn = x.shape[1]
    if strategy == "uniform":
        if value_range is not None:
            lo = torch.full((Fn,), float(value_range[0]), device=x.device)
            hi = torch.full((Fn,), float(value_range[1]), device=x.device)
        else:
            lo, hi = x.min(dim=0).values, x.max(dim=0).values
        steps = torch.arange(1, n_bits + 1, device=x.device, dtype=torch.float32) / (n_bits + 1)
        return lo.unsqueeze(1) + (hi - lo).unsqueeze(1) * steps.unsqueeze(0)
    if strategy == "quantile":
        q = torch.arange(1, n_bits + 1, device=x.device, dtype=torch.float32) / (n_bits + 1)
        return torch.quantile(x, q, dim=0).transpose(0, 1).contiguous()
    if strategy == "unique":
        out = torch.empty(Fn, n_bits, device=x.device)
        for f in range(Fn):
            u = torch.unique(x[:, f])
            if u.numel() > 1:
                u = u[1:]  # ``x >= min`` is always True — drop it
            if u.numel() >= n_bits:
                idx = torch.linspace(0, u.numel() - 1, n_bits, device=x.device).round().long()
                out[f] = u[idx]
            else:  # pad by repeating the max (bits that are never on for training data)
                pad = u.max().repeat(n_bits - u.numel()) if u.numel() else torch.zeros(n_bits)
                out[f] = torch.cat([u, pad.to(u.device)])
        return out
    raise ValueError("strategy must be 'quantile', 'uniform' or 'unique'")

Graphs

torchtsetlin.data.graphs

Graph-structured input for :class:~torchtsetlin.models.GraphTsetlinMachine.

A :class:GraphBatch holds a batch of directed multigraphs whose nodes carry Boolean feature vectors and whose edges carry an integer type. Node features are usually sparse hypervectors of symbols, built by a :class:SymbolTable (Halenka et al., 2024; Granmo et al., 2025, arXiv:2507.14874), but any Boolean vector works: an image patch with its position bits is a node, so a graph without edges is the input of a convolutional Tsetlin machine.

Storage is padded: features (B, P, F) with P the largest node count of the batch and a (B, P) node mask. Edges are kept per graph, sorted by graph, with local node indices, so that slicing a batch (batch[idx]) is cheap and keeps every edge inside its graph.

SymbolTable

SymbolTable(symbols: Iterable, dim: int = 128, bits: int = 2, encoding: str = 'random', seed: int = 0)

Map symbol names to sparse hypervector positions (the node properties of a graph).

Every symbol owns bits positions out of dim. A node's feature vector is the OR (bundle) of its symbols' positions. The three encodings follow the reference GraphTsetlinMachine (Graphs in graphs.py):

• "random": bits distinct random positions per symbol (seeded);
• "double_hashing": positions i mod dim and p - (i mod p) with p the largest prime below dim (always 2 bits);
• "one_hot": dim = len(symbols) and symbol i owns position i. This is a plain Boolean feature per symbol, and it makes decoding exact.

Parameters:

Name	Type	Description	Default
symbols	Iterable	the symbol names (any hashable values).	required
dim	int	hypervector size (ignored for "one_hot").	128
bits	int	positions per symbol (ignored for "double_hashing" and "one_hot").	2
encoding	str	"random", "double_hashing" or "one_hot".	'random'
seed	int	seed of the random positions.	0

Source code in src/torchtsetlin/data/graphs.py

def __init__(
    self,
    symbols: Iterable,
    dim: int = 128,
    bits: int = 2,
    encoding: str = "random",
    seed: int = 0,
) -> None:
    self.symbols = list(symbols)
    if len(set(self.symbols)) != len(self.symbols):
        raise ValueError("symbol names must be unique")
    self.index = {s: i for i, s in enumerate(self.symbols)}
    n = len(self.symbols)
    if encoding == "one_hot":
        dim, bits = n, 1
        pos = torch.arange(n).unsqueeze(1)
    elif encoding == "double_hashing":
        bits = 2
        p = _prev_prime(dim)
        i = torch.arange(n)
        pos = torch.stack([i % dim, p - (i % p)], dim=1)
    elif encoding == "random":
        if not 1 <= bits <= dim:
            raise ValueError("bits must lie in [1, dim]")
        g = torch.Generator().manual_seed(int(seed))
        pos = torch.stack([torch.randperm(dim, generator=g)[:bits] for _ in range(n)]) if n else \
            torch.zeros(0, bits, dtype=torch.long)
    else:
        raise ValueError("encoding must be 'random', 'double_hashing' or 'one_hot'")
    self.dim = int(dim)
    self.bits = int(bits)
    self.encoding = encoding
    self.positions: Tensor = pos.long()  # (n_symbols, bits)

vector

vector(symbols: Iterable) -> Tensor

(dim,) Boolean hypervector of a set of symbols (their bundle).

Source code in src/torchtsetlin/data/graphs.py

def vector(self, symbols: Iterable) -> Tensor:
    """``(dim,)`` Boolean hypervector of a set of symbols (their bundle)."""
    out = torch.zeros(self.dim, dtype=torch.bool)
    for s in symbols:
        out[self.positions[self.index[s]]] = True
    return out

encode

encode(nodes: Sequence[Iterable]) -> Tensor

(n, dim) Boolean feature matrix, one row per node's symbol set.

Source code in src/torchtsetlin/data/graphs.py

def encode(self, nodes: Sequence[Iterable]) -> Tensor:
    """``(n, dim)`` Boolean feature matrix, one row per node's symbol set."""
    if not len(nodes):
        return torch.zeros(0, self.dim, dtype=torch.bool)
    return torch.stack([self.vector(s) for s in nodes])

decode

decode(vector: Tensor) -> List

Symbols whose positions are all set in vector ((dim,) Boolean).

Source code in src/torchtsetlin/data/graphs.py

def decode(self, vector: Tensor) -> List:
    """Symbols whose positions are all set in ``vector`` (``(dim,)`` Boolean)."""
    v = as_bool_tensor(vector).flatten()
    hit = v[self.positions].all(dim=1)
    return [self.symbols[i] for i in torch.nonzero(hit).flatten().tolist()]

feature_names

feature_names() -> List[str]

Names of the dim positions: the symbol for one-hot tables, else hv[k].

Source code in src/torchtsetlin/data/graphs.py

def feature_names(self) -> List[str]:
    """Names of the ``dim`` positions: the symbol for one-hot tables, else ``hv[k]``."""
    if self.encoding == "one_hot":
        return [str(s) for s in self.symbols]
    return [f"hv[{k}]" for k in range(self.dim)]

GraphBatch

GraphBatch(node_features: Tensor, n_nodes: Tensor, edges: Optional[Tensor] = None, edge_type: Optional[Tensor] = None, edges_per_graph: Optional[Tensor] = None, node_type: Optional[Tensor] = None, n_edge_types: Optional[int] = None, n_node_types: Optional[int] = None, edge_type_names: Optional[Sequence[str]] = None)

A padded batch of directed multigraphs with typed edges.

Parameters:

Name	Type	Description	Default
node_features	Tensor	(B, P, F) Boolean node features; rows beyond a graph's node count are padding and are ignored.	required
n_nodes	Tensor	(B,) node count per graph.	required
edges	Optional[Tensor]	(E, 2) long (source, destination) pairs of local node indices, grouped by graph in batch order.	None
edge_type	Optional[Tensor]	(E,) integer edge types in [0, n_edge_types).	None
edges_per_graph	Optional[Tensor]	(B,) number of edges of each graph (they sum to E).	None
node_type	Optional[Tensor]	optional (B, P) integer node types in [0, n_node_types).	None
n_edge_types	Optional[int]	number of edge types of the dataset (kept by slices, so a batch that happens to lack a type still declares it). Defaults to max(edge_type) + 1.	None
n_node_types	Optional[int]	number of node types of the dataset. Defaults to max(node_type) + 1.	None
edge_type_names	Optional[Sequence[str]]	optional names of the edge types, for printing clauses.	None

Messages travel along an edge from its source to its destination, so a node learns about the nodes that have an edge into it. Undirected relations need an edge each way, usually with two types (as in the sequence demo: "left" and "right").

Build batches with :meth:from_graphs, :meth:from_sequences or :meth:from_images; index them like tensors along the graph axis (batch[3], batch[2:10], batch[torch.tensor([4, 1])]) and concatenate them with :meth:cat.

Source code in src/torchtsetlin/data/graphs.py

def __init__(
    self,
    node_features: Tensor,
    n_nodes: Tensor,
    edges: Optional[Tensor] = None,
    edge_type: Optional[Tensor] = None,
    edges_per_graph: Optional[Tensor] = None,
    node_type: Optional[Tensor] = None,
    n_edge_types: Optional[int] = None,
    n_node_types: Optional[int] = None,
    edge_type_names: Optional[Sequence[str]] = None,
) -> None:
    x = as_bool_tensor(node_features)
    if x.dim() != 3:
        raise ValueError("node_features must have shape (B, P, F)")
    dev = x.device
    B, P, _ = x.shape
    n_nodes = torch.as_tensor(n_nodes, dtype=torch.long, device=dev).reshape(-1)
    if n_nodes.numel() != B:
        raise ValueError("n_nodes must hold one count per graph")
    if B and (int(n_nodes.min()) < 0 or int(n_nodes.max()) > P):
        raise ValueError("node counts must lie in [0, P]")
    if edges is None:
        edges = torch.zeros(0, 2, dtype=torch.long, device=dev)
    edges = torch.as_tensor(edges, dtype=torch.long, device=dev).reshape(-1, 2)
    E = edges.shape[0]
    edge_type = (torch.zeros(E, dtype=torch.long, device=dev) if edge_type is None
                 else torch.as_tensor(edge_type, dtype=torch.long, device=dev).reshape(-1))
    if edge_type.numel() != E:
        raise ValueError("edge_type must hold one type per edge")
    if edges_per_graph is None:
        if E and B != 1:
            raise ValueError("edges_per_graph is required for a batch of several graphs")
        edges_per_graph = torch.full((B,), E, dtype=torch.long, device=dev)
    edges_per_graph = torch.as_tensor(edges_per_graph, dtype=torch.long, device=dev).reshape(-1)
    if edges_per_graph.numel() != B or int(edges_per_graph.sum()) != E:
        raise ValueError("edges_per_graph must hold one count per graph and sum to E")
    if E:
        graph = torch.repeat_interleave(torch.arange(B, device=dev), edges_per_graph)
        if int(edges.min()) < 0 or bool((edges >= n_nodes[graph].unsqueeze(1)).any()):
            raise ValueError("an edge refers to a node outside its graph")
        if int(edge_type.min()) < 0:
            raise ValueError("edge types must be non-negative")
    if node_type is not None:
        node_type = torch.as_tensor(node_type, dtype=torch.long, device=dev)
        if node_type.shape != (B, P):
            raise ValueError("node_type must have shape (B, P)")
    self.node_features = x
    self.n_nodes = n_nodes
    self.edges = edges
    self.edge_type = edge_type
    self.edges_per_graph = edges_per_graph
    self.node_type = node_type
    inferred_et = int(edge_type.max()) + 1 if E else 1
    self.n_edge_types = int(n_edge_types) if n_edge_types is not None else inferred_et
    if E and inferred_et > self.n_edge_types:
        raise ValueError("an edge type is >= n_edge_types")
    inferred_nt = int(node_type.max()) + 1 if node_type is not None and node_type.numel() else 1
    self.n_node_types = int(n_node_types) if n_node_types is not None else inferred_nt
    if inferred_nt > self.n_node_types:
        raise ValueError("a node type is >= n_node_types")
    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")

shape property

shape: Size

(B, P, F), the shape of the padded node features.

node_mask property

node_mask: Tensor

(B, P) Boolean: which rows are real nodes.

edge_graph property

edge_graph: Tensor

(E,) graph index of every edge.

cat staticmethod

cat(batches: Sequence[GraphBatch]) -> GraphBatch

Concatenate batches along the graph axis (padding to the largest node count).

Source code in src/torchtsetlin/data/graphs.py

@staticmethod
def cat(batches: Sequence[GraphBatch]) -> GraphBatch:
    """Concatenate batches along the graph axis (padding to the largest node count)."""
    if not batches:
        raise ValueError("nothing to concatenate")
    P = max(b.shape[1] for b in batches)
    F_ = batches[0].n_features
    if any(b.n_features != F_ for b in batches):
        raise ValueError("all batches need the same feature count")
    dev = batches[0].device
    pad = lambda t, v: torch.nn.functional.pad(t, (0, 0, 0, P - t.shape[1]), value=v) if t.dim() == 3 \
        else torch.nn.functional.pad(t, (0, P - t.shape[1]), value=v)  # noqa: E731
    feats = torch.cat([pad(b.node_features.to(dev), False) for b in batches])
    has_type = any(b.node_type is not None for b in batches)
    ntype = None
    if has_type:
        ntype = torch.cat([pad(b.node_type.to(dev) if b.node_type is not None else
                               torch.zeros(b.shape[:2], dtype=torch.long, device=dev), 0) for b in batches])
    return GraphBatch(
        feats, torch.cat([b.n_nodes.to(dev) for b in batches]),
        torch.cat([b.edges.to(dev) for b in batches]), torch.cat([b.edge_type.to(dev) for b in batches]),
        torch.cat([b.edges_per_graph.to(dev) for b in batches]), ntype,
        max(b.n_edge_types for b in batches), max(b.n_node_types for b in batches),
        batches[0].edge_type_names)

flat_edges

flat_edges() -> Tuple[Tensor, Tensor, Tensor]

(src, dst, type) with node indices into the flattened B * P node axis.

Source code in src/torchtsetlin/data/graphs.py

def flat_edges(self) -> Tuple[Tensor, Tensor, Tensor]:
    """``(src, dst, type)`` with node indices into the flattened ``B * P`` node axis."""
    P = self.node_features.shape[1]
    base = self.edge_graph * P
    return base + self.edges[:, 0], base + self.edges[:, 1], self.edge_type

from_graphs classmethod

from_graphs(node_features: Sequence, edges: Optional[Sequence] = None, edge_types: Optional[Sequence] = None, node_types: Optional[Sequence] = None, n_edge_types: Optional[int] = None, n_node_types: Optional[int] = None, edge_type_names: Optional[Sequence[str]] = None) -> GraphBatch

Build a batch from per-graph lists.

Parameters:

Name	Type	Description	Default
node_features	Sequence	per graph, an (n_i, F) Boolean array.	required
edges	Optional[Sequence]	per graph, an (E_i, 2) array of (source, destination) local node indices (or None / empty for no edges).	None
edge_types	Optional[Sequence]	per graph, an (E_i,) array of edge types (default: all 0).	None
node_types	Optional[Sequence]	per graph, an (n_i,) array of node types (default: none).	None

Source code in src/torchtsetlin/data/graphs.py

@classmethod
def from_graphs(
    cls,
    node_features: Sequence,
    edges: Optional[Sequence] = None,
    edge_types: Optional[Sequence] = None,
    node_types: Optional[Sequence] = None,
    n_edge_types: Optional[int] = None,
    n_node_types: Optional[int] = None,
    edge_type_names: Optional[Sequence[str]] = None,
) -> GraphBatch:
    """Build a batch from per-graph lists.

    Args:
        node_features: per graph, an ``(n_i, F)`` Boolean array.
        edges: per graph, an ``(E_i, 2)`` array of ``(source, destination)`` local node
            indices (or ``None`` / empty for no edges).
        edge_types: per graph, an ``(E_i,)`` array of edge types (default: all ``0``).
        node_types: per graph, an ``(n_i,)`` array of node types (default: none).
    """
    feats = [as_bool_tensor(f) for f in node_features]
    if not feats:
        raise ValueError("no graphs")
    F_ = feats[0].shape[1] if feats[0].dim() == 2 else 0
    n = torch.tensor([f.shape[0] for f in feats], dtype=torch.long)
    P = max(1, int(n.max()))
    X = torch.zeros(len(feats), P, F_, dtype=torch.bool)
    for i, f in enumerate(feats):
        if f.dim() != 2 or f.shape[1] != F_:
            raise ValueError("every graph needs (n_i, F) node features with the same F")
        X[i, : f.shape[0]] = f.cpu()
    E_list, T_list, per = [], [], []
    for i in range(len(feats)):
        e = edges[i] if edges is not None else None
        e = torch.zeros(0, 2, dtype=torch.long) if e is None else torch.as_tensor(e, dtype=torch.long).reshape(-1, 2)
        t = edge_types[i] if edge_types is not None else None
        t = torch.zeros(e.shape[0], dtype=torch.long) if t is None else torch.as_tensor(t, dtype=torch.long).reshape(-1)
        E_list.append(e)
        T_list.append(t)
        per.append(e.shape[0])
    nt = None
    if node_types is not None:
        nt = torch.zeros(len(feats), P, dtype=torch.long)
        for i, t in enumerate(node_types):
            t = torch.as_tensor(t, dtype=torch.long).reshape(-1)
            nt[i, : t.numel()] = t
    return cls(X, n, torch.cat(E_list), torch.cat(T_list), torch.tensor(per, dtype=torch.long), nt,
               n_edge_types, n_node_types, edge_type_names)

from_sequences classmethod

from_sequences(sequences: Sequence[Sequence], symbols: SymbolTable, edge_types: Tuple[str, str] = ('left', 'right'), unknown: str = 'ignore') -> GraphBatch

One node per sequence element, chained in both directions.

Node k sends to k - 1 over edge type 0 (named edge_types[0]) and to k + 1 over type 1, as in the reference SequenceClassificationDemo. An element may be a single symbol or a set of symbols. Symbols missing from the table are ignored (unknown="ignore") or rejected ("error").

Source code in src/torchtsetlin/data/graphs.py

@classmethod
def from_sequences(
    cls,
    sequences: Sequence[Sequence],
    symbols: SymbolTable,
    edge_types: Tuple[str, str] = ("left", "right"),
    unknown: str = "ignore",
) -> GraphBatch:
    """One node per sequence element, chained in both directions.

    Node ``k`` sends to ``k - 1`` over edge type 0 (named ``edge_types[0]``) and to
    ``k + 1`` over type 1, as in the reference ``SequenceClassificationDemo``. An element
    may be a single symbol or a set of symbols. Symbols missing from the table are
    ignored (``unknown="ignore"``) or rejected (``"error"``).
    """
    if unknown not in ("ignore", "error"):
        raise ValueError("unknown must be 'ignore' or 'error'")
    feats, edges, types = [], [], []
    for seq in sequences:
        nodes = []
        for el in seq:
            group = el if isinstance(el, (set, frozenset, list, tuple)) else [el]
            kept = [s for s in group if s in symbols.index]
            if unknown == "error" and len(kept) != len(group):
                raise KeyError(f"unknown symbol in {group!r}")
            nodes.append(kept)
        feats.append(symbols.encode(nodes) if nodes else torch.zeros(0, symbols.dim, dtype=torch.bool))
        L = len(nodes)
        k = torch.arange(L)
        left = torch.stack([k[1:], k[:-1]], dim=1)   # k -> k-1, type 0
        right = torch.stack([k[:-1], k[1:]], dim=1)  # k -> k+1, type 1
        edges.append(torch.cat([left, right]))
        types.append(torch.cat([torch.zeros(max(L - 1, 0), dtype=torch.long),
                                torch.ones(max(L - 1, 0), dtype=torch.long)]))
    return cls.from_graphs(feats, edges, types, n_edge_types=2, edge_type_names=edge_types)

from_images classmethod

from_images(images, patch_size: Union[int, Tuple[int, int]], stride: Union[int, Tuple[int, int]] = 1, position_encoding: bool = True) -> GraphBatch

One node per image patch and no edges (paper Sec. 3.2, "disconnected nodes").

Node features are the Z * kh * kw patch pixels followed by thermometer-coded patch coordinates, exactly the patch features of :class:~torchtsetlin.models.ConvTsetlinMachine with the same arguments. A depth-1 graph machine on these graphs is therefore the convolutional coalesced machine.

Parameters:

Name	Type	Description	Default
images		Boolean (B, Z, H, W) (or (B, H, W)) images.	required

Source code in src/torchtsetlin/data/graphs.py

@classmethod
def from_images(
    cls,
    images,
    patch_size: Union[int, Tuple[int, int]],
    stride: Union[int, Tuple[int, int]] = 1,
    position_encoding: bool = True,
) -> GraphBatch:
    """One node per image patch and no edges (paper Sec. 3.2, "disconnected nodes").

    Node features are the ``Z * kh * kw`` patch pixels followed by thermometer-coded
    patch coordinates, exactly the patch features of
    :class:`~torchtsetlin.models.ConvTsetlinMachine` with the same arguments. A depth-1
    graph machine on these graphs is therefore the convolutional coalesced machine.

    Args:
        images: Boolean ``(B, Z, H, W)`` (or ``(B, H, W)``) images.
    """
    x = as_bool_tensor(images)
    if x.dim() == 3:
        x = x.unsqueeze(1)
    if x.dim() != 4:
        raise ValueError("images must have shape (B, Z, H, W) or (B, H, W)")
    kh, kw = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size
    sh, sw = (stride, stride) if isinstance(stride, int) else stride
    B, Z, H, W = x.shape
    Py, Px = (H - kh) // sh + 1, (W - kw) // sw + 1
    if Py < 1 or Px < 1:
        raise ValueError("image smaller than the patch")
    patches = torch.nn.functional.unfold(x.float(), kernel_size=(kh, kw), stride=(sh, sw))
    patches = patches.transpose(1, 2) > 0.5  # (B, P, Z*kh*kw)
    if position_encoding:
        ys = torch.arange(Py, device=x.device).repeat_interleave(Px)
        xs = torch.arange(Px, device=x.device).repeat(Py)
        pos = torch.cat([ys.unsqueeze(1) > torch.arange(Py - 1, device=x.device).unsqueeze(0),
                         xs.unsqueeze(1) > torch.arange(Px - 1, device=x.device).unsqueeze(0)], dim=1)
        patches = torch.cat([patches, pos.unsqueeze(0).expand(B, -1, -1)], dim=2)
    P = Py * Px
    n = torch.full((B,), P, dtype=torch.long, device=x.device)
    return cls(patches, n, None, None, torch.zeros(B, dtype=torch.long, device=x.device), None, 1, 1)

Datasets

torchtsetlin.data.datasets

Ready-made Boolean datasets and helpers to wrap your own data.

BooleanTensorDataset

BooleanTensorDataset(x, y, device=None, regression: bool = False)

Bases: TensorDataset

TensorDataset that coerces features to torch.bool and targets to long (or float for regression). Keeps everything on one device for fast GPU epochs.

Source code in src/torchtsetlin/data/datasets.py

def __init__(self, x, y, device=None, regression: bool = False) -> None:
    xb = as_bool_tensor(x, device=torch.device(device) if device else None)
    if regression:
        yt = torch.as_tensor(np.asarray(y), dtype=torch.float32, device=xb.device)
    else:
        yt = as_long_tensor(y, device=xb.device)
    super().__init__(xb, yt)

TransformDataset

TransformDataset(base: Dataset, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None)

Bases: Dataset

Apply a (booleanizing) transform lazily to the samples of another dataset.

Source code in src/torchtsetlin/data/datasets.py

def __init__(self, base: Dataset, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None):
    self.base = base
    self.transform = transform
    self.target_transform = target_transform

to_boolean_tensors

to_boolean_tensors(x, y=None, device=None) -> Tuple[Tensor, Optional[Tensor]]

Coerce x to a Boolean tensor and y to a long tensor on device.

Source code in src/torchtsetlin/data/datasets.py

def to_boolean_tensors(x, y=None, device=None) -> Tuple[Tensor, Optional[Tensor]]:
    """Coerce ``x`` to a Boolean tensor and ``y`` to a long tensor on ``device``."""
    dev = torch.device(device) if device else None
    xb = as_bool_tensor(x, device=dev)
    yb = as_long_tensor(y, device=dev) if y is not None else None
    return xb, yb

make_xor

make_xor(n: int, n_features: int = 2, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]

Random Boolean vectors labelled by x0 XOR x1 (other features are distractors).

Source code in src/torchtsetlin/data/datasets.py

def make_xor(n: int, n_features: int = 2, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]:
    """Random Boolean vectors labelled by ``x0 XOR x1`` (other features are distractors)."""
    return make_noisy_xor(n, n_features=n_features, noise=0.0, seed=seed, device=device)

make_noisy_xor

make_noisy_xor(n: int, n_features: int = 12, noise: float = 0.4, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]

The Noisy XOR benchmark of the original Tsetlin machine paper: 12 Boolean features of which only the first two matter (y = x0 XOR x1); a fraction noise of the training labels are flipped.

Source code in src/torchtsetlin/data/datasets.py

def make_noisy_xor(
    n: int, n_features: int = 12, noise: float = 0.4, seed: Optional[int] = 0, device=None
) -> Tuple[Tensor, Tensor]:
    """The *Noisy XOR* benchmark of the original Tsetlin machine paper: 12 Boolean features
    of which only the first two matter (``y = x0 XOR x1``); a fraction ``noise`` of the
    training labels are flipped."""
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    x = torch.randint(0, 2, (n, n_features), generator=g).bool()
    y = (x[:, 0] ^ x[:, 1]).long()
    if noise > 0:
        flip = torch.rand(n, generator=g) < noise
        y = torch.where(flip, 1 - y, y)
    dev = torch.device(device) if device else None
    return (x.to(dev), y.to(dev)) if dev else (x, y)

make_parity

make_parity(n: int, n_bits: int = 4, n_features: int = 8, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]

Parity of the first n_bits features (harder than XOR).

Source code in src/torchtsetlin/data/datasets.py

def make_parity(n: int, n_bits: int = 4, n_features: int = 8, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]:
    """Parity of the first ``n_bits`` features (harder than XOR)."""
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    x = torch.randint(0, 2, (n, n_features), generator=g).bool()
    y = (x[:, :n_bits].long().sum(dim=1) % 2).long()
    dev = torch.device(device) if device else None
    return (x.to(dev), y.to(dev)) if dev else (x, y)

make_2d_noisy_xor

make_2d_noisy_xor(n: int, size: int = 4, noise: float = 0.0, seed: Optional[int] = 0, device=None, max_tries: int = 200) -> Tuple[Tensor, Tensor]

2D Noisy XOR of the Convolutional TM paper: a size x size binary image whose pixels are random, except for one 2 x 2 patch (at a random position) that carries the class pattern.

Class 1 patterns are the diagonals [[1,0],[0,1]] / [[0,1],[1,0]], class 0 patterns are the horizontal bars [[1,1],[0,0]] / [[0,0],[1,1]]. Backgrounds are resampled so that no other window contains any of the four patterns (otherwise the label would be ambiguous). A fraction noise of the labels is flipped.

Returns images (n, 1, size, size) and labels (n,).

Source code in src/torchtsetlin/data/datasets.py

def make_2d_noisy_xor(
    n: int, size: int = 4, noise: float = 0.0, seed: Optional[int] = 0, device=None, max_tries: int = 200
) -> Tuple[Tensor, Tensor]:
    """2D Noisy XOR of the Convolutional TM paper: a ``size x size`` binary image whose
    pixels are random, except for one ``2 x 2`` patch (at a random position) that carries
    the class pattern.

    Class ``1`` patterns are the diagonals ``[[1,0],[0,1]]`` / ``[[0,1],[1,0]]``, class ``0``
    patterns are the horizontal bars ``[[1,1],[0,0]]`` / ``[[0,0],[1,1]]``. Backgrounds are
    resampled so that no *other* window contains any of the four patterns (otherwise the
    label would be ambiguous). A fraction ``noise`` of the labels is flipped.

    Returns images ``(n, 1, size, size)`` and labels ``(n,)``.
    """
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    patterns = torch.tensor(
        [[[1, 1], [0, 0]], [[0, 0], [1, 1]], [[1, 0], [0, 1]], [[0, 1], [1, 0]]], dtype=torch.bool
    )  # first two -> class 0, last two -> class 1
    y = torch.randint(0, 2, (n,), generator=g)
    variant = torch.randint(0, 2, (n,), generator=g)
    which = y * 2 + variant
    py = torch.randint(0, size - 1, (n,), generator=g)
    px = torch.randint(0, size - 1, (n,), generator=g)
    x = torch.randint(0, 2, (n, 1, size, size), generator=g).bool()
    ar = torch.arange(n)

    def stamp(imgs: Tensor, idx: Tensor) -> None:
        pat = patterns[which[idx]]  # (k, 2, 2)
        for dy in range(2):
            for dx in range(2):
                imgs[idx, 0, py[idx] + dy, px[idx] + dx] = pat[:, dy, dx]

    def ambiguous(imgs: Tensor) -> Tensor:
        win = imgs[:, 0].unfold(1, 2, 1).unfold(2, 2, 1)  # (n, size-1, size-1, 2, 2)
        hit = torch.zeros(imgs.shape[0], size - 1, size - 1, dtype=torch.bool)
        for pat in patterns:
            hit |= (win == pat).all(dim=-1).all(dim=-1)
        hit[ar, py, px] = False  # the placed patch is allowed
        return hit.any(dim=(1, 2))

    stamp(x, ar)
    bad = ambiguous(x)
    tries = 0
    while bool(bad.any()) and tries < max_tries:
        idx = torch.nonzero(bad).flatten()
        x[idx] = torch.randint(0, 2, (idx.numel(), 1, size, size), generator=g).bool()
        stamp(x, idx)
        bad = ambiguous(x)
        tries += 1
    if noise > 0:
        flip = torch.rand(n, generator=g) < noise
        y = torch.where(flip, 1 - y, y)
    dev = torch.device(device) if device else None
    return (x.to(dev), y.long().to(dev)) if dev else (x, y.long())

make_shapes

make_shapes(n: int, size: int = 8, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]

Toy image classification: a hollow 3x3 circle (class 0) or a 3x3 cross (class 1) at a random position in a size x size image (Chapter 4 of the Tsetlin machine book). Returns (n, 1, size, size) Boolean images and labels.

Source code in src/torchtsetlin/data/datasets.py

def make_shapes(n: int, size: int = 8, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]:
    """Toy image classification: a hollow ``3x3`` circle (class 0) or a ``3x3`` cross (class 1)
    at a random position in a ``size x size`` image (Chapter 4 of the Tsetlin machine book).
    Returns ``(n, 1, size, size)`` Boolean images and labels."""
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    x = torch.zeros(n, 1, size, size, dtype=torch.bool)
    y = torch.randint(0, 2, (n,), generator=g)
    circle = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 1, 0]], dtype=torch.bool)
    cross = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=torch.bool)
    py = torch.randint(0, size - 2, (n,), generator=g)
    px = torch.randint(0, size - 2, (n,), generator=g)
    for i in range(n):
        x[i, 0, py[i] : py[i] + 3, px[i] : px[i] + 3] = cross if y[i] == 1 else circle
    dev = torch.device(device) if device else None
    return (x.to(dev), y.long().to(dev)) if dev else (x, y.long())

make_3d_shapes

make_3d_shapes(n: int, size: Tuple[int, int, int] = (8, 12, 12), seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]

Toy volume classification that needs depth context: a 3x3x3 shape at a random position in a (D, H, W) volume.

Class 0 (spindle) stacks the 2D slices plus, dot, plus; class 1 (3D cross) stacks dot, plus, dot. Every single slice therefore looks the same in both classes (a plus or a dot) — only their order along depth tells them apart, which a 3D window with kd = 3 sees in one clause.

Returns (n, 1, D, H, W) Boolean volumes and labels.

Source code in src/torchtsetlin/data/datasets.py

def make_3d_shapes(
    n: int, size: Tuple[int, int, int] = (8, 12, 12), seed: Optional[int] = 0, device=None
) -> Tuple[Tensor, Tensor]:
    """Toy volume classification that needs depth context: a ``3x3x3`` shape at a random
    position in a ``(D, H, W)`` volume.

    Class 0 (*spindle*) stacks the 2D slices ``plus, dot, plus``; class 1 (*3D cross*) stacks
    ``dot, plus, dot``. Every single slice therefore looks the same in both classes (a plus
    or a dot) — only their order along depth tells them apart, which a 3D window with
    ``kd = 3`` sees in one clause.

    Returns ``(n, 1, D, H, W)`` Boolean volumes and labels."""
    D, H, W = (int(v) for v in size)
    if min(D, H, W) < 3:
        raise ValueError("every side of the volume must be at least 3")
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    x = torch.zeros(n, 1, D, H, W, dtype=torch.bool)
    y = torch.randint(0, 2, (n,), generator=g)
    plus = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=torch.bool)
    dot = torch.tensor([[0, 0, 0], [0, 1, 0], [0, 0, 0]], dtype=torch.bool)
    spindle = torch.stack([plus, dot, plus])
    cross3d = torch.stack([dot, plus, dot])
    pd = torch.randint(0, D - 2, (n,), generator=g)
    py = torch.randint(0, H - 2, (n,), generator=g)
    px = torch.randint(0, W - 2, (n,), generator=g)
    for i in range(n):
        x[i, 0, pd[i] : pd[i] + 3, py[i] : py[i] + 3, px[i] : px[i] + 3] = (
            cross3d if y[i] == 1 else spindle
        )
    dev = torch.device(device) if device else None
    return (x.to(dev), y.long().to(dev)) if dev else (x, y.long())

make_segmentation_shapes

make_segmentation_shapes(n: int, size: int = 16, seed: Optional[int] = 0, device=None) -> Tuple[Tensor, Tensor]

Toy dense-prediction task: label every pixel background / disc / ring.

Two solid discs and two rings of radius 2-3 are stamped into a size x size Boolean image. Disc interiors and ring interiors look identical on their own, so the classes are only separable from a pixel's neighbourhood — which is exactly what a dense Tsetlin machine gets to see, and what a per-pixel lookup cannot solve.

Returns (n, 1, size, size) Boolean images and (n, size, size) long label maps with classes 0 = background, 1 = disc, 2 = ring.

Source code in src/torchtsetlin/data/datasets.py

def make_segmentation_shapes(
    n: int, size: int = 16, seed: Optional[int] = 0, device=None
) -> Tuple[Tensor, Tensor]:
    """Toy dense-prediction task: label every pixel ``background`` / ``disc`` / ``ring``.

    Two solid discs and two rings of radius 2-3 are stamped into a ``size x size`` Boolean
    image. Disc interiors and ring interiors look identical on their own, so the classes are
    only separable from a pixel's *neighbourhood* — which is exactly what a dense Tsetlin
    machine gets to see, and what a per-pixel lookup cannot solve.

    Returns ``(n, 1, size, size)`` Boolean images and ``(n, size, size)`` long label maps
    with classes ``0 = background, 1 = disc, 2 = ring``.
    """
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    x = torch.zeros(n, 1, size, size, dtype=torch.bool)
    y = torch.zeros(n, size, size, dtype=torch.long)
    yy, xx = torch.meshgrid(torch.arange(size), torch.arange(size), indexing="ij")
    for i in range(n):
        for cls in (1, 2):
            for _ in range(2):
                r = int(torch.randint(2, 4, (1,), generator=g))
                cy = int(torch.randint(r + 1, size - r - 1, (1,), generator=g))
                cx = int(torch.randint(r + 1, size - r - 1, (1,), generator=g))
                d2 = (yy - cy) ** 2 + (xx - cx) ** 2
                solid = d2 <= r * r
                shape = solid if cls == 1 else (solid & (d2 > (r - 1) ** 2))
                x[i, 0] |= shape
                y[i][solid] = cls
    dev = torch.device(device) if device else None
    return (x.to(dev), y.to(dev)) if dev else (x, y)

make_scenes

make_scenes(n: int, size: Union[int, Tuple[int, int]] = 32, seed: Optional[int] = 0, device: Optional[Union[str, device]] = None) -> Tuple[Tensor, Tensor]

Synthetic CamVid-like street scenes for semantic segmentation.

Generates the properties that make CamVid hard rather than the photographs themselves (which are not redistributable): a wavy horizon, a large contiguous sky, buildings standing on the horizon, trees, and a road, with per-class texture (window grids, foliage noise, lane markings) and a global illumination jitter so that the classes are not separable by colour alone. Class frequencies are heavily imbalanced, as in CamVid.

Parameters:

Name	Type	Description	Default
n	int	number of scenes.	required
size	Union[int, Tuple[int, int]]	H or (H, W).	32
seed	Optional[int]	RNG seed (None = unseeded).	0
device	Optional[Union[str, device]]	device for the returned tensors.	None

Returns:

Type	Description
Tensor	(rgb, labels) with rgb float (n, 3, H, W) in [0, 1] and labels
Tensor	long (n, H, W) indexing :data:SCENE_CLASSES
Tuple[Tensor, Tensor]	(0 = sky, 1 = building, 2 = tree, 3 = road).

Booleanize the images before handing them to a model, for example with :class:~torchtsetlin.data.ColorThermometerEncoder, which keeps the ordering between intensity levels that a plain threshold would throw away.

Source code in src/torchtsetlin/data/datasets.py

def make_scenes(
    n: int,
    size: Union[int, Tuple[int, int]] = 32,
    seed: Optional[int] = 0,
    device: Optional[Union[str, torch.device]] = None,
) -> Tuple[Tensor, Tensor]:
    """Synthetic CamVid-like street scenes for semantic segmentation.

    Generates the properties that make CamVid hard rather than the photographs themselves
    (which are not redistributable): a wavy **horizon**, a large contiguous **sky**,
    **buildings** standing on the horizon, **trees**, and a **road**, with per-class texture
    (window grids, foliage noise, lane markings) and a global illumination jitter so that the
    classes are not separable by colour alone. Class frequencies are heavily imbalanced, as
    in CamVid.

    Args:
        n: number of scenes.
        size: ``H`` or ``(H, W)``.
        seed: RNG seed (``None`` = unseeded).
        device: device for the returned tensors.

    Returns:
        ``(rgb, labels)`` with ``rgb`` float ``(n, 3, H, W)`` in ``[0, 1]`` and ``labels``
        long ``(n, H, W)`` indexing :data:`SCENE_CLASSES`
        (``0 = sky, 1 = building, 2 = tree, 3 = road``).

    Booleanize the images before handing them to a model, for example with
    :class:`~torchtsetlin.data.ColorThermometerEncoder`, which keeps the ordering between
    intensity levels that a plain threshold would throw away.
    """
    SKY, BUILDING, TREE, ROAD = 0, 1, 2, 3
    H, W = (size, size) if isinstance(size, int) else (int(size[0]), int(size[1]))
    if H < 12 or W < 12:
        raise ValueError("scenes need to be at least 12x12")
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    rgb = torch.zeros(n, 3, H, W)
    lab = torch.full((n, H, W), ROAD, dtype=torch.long)
    xs = torch.arange(W).float()
    rows = torch.arange(H).unsqueeze(1)
    base = {
        SKY: (120, 160, 215),
        BUILDING: (150, 145, 135),
        TREE: (70, 120, 65),
        ROAD: (95, 95, 100),
    }
    two_pi = 2 * math.pi

    def rint(lo: int, hi: int) -> int:
        return int(torch.randint(lo, max(lo + 1, hi), (1,), generator=g))

    for i in range(n):
        h0 = rint(H // 3, H // 2)
        amp = float(torch.rand(1, generator=g)) * 2.0
        phase = float(torch.rand(1, generator=g)) * two_pi
        horizon = (h0 + amp * torch.sin(xs / W * two_pi + phase)).round().long().clamp(2, H - 4)
        lab[i] = torch.where(rows < horizon.unsqueeze(0), SKY, ROAD)

        for cls, count, wlo, whi, hlo, hhi in (
            (BUILDING, rint(1, 4), 3, 9, 4, 12),
            (TREE, rint(0, 3), 3, 6, 4, 9),
        ):
            for _ in range(count):
                bw = rint(wlo, whi)
                bx = rint(0, max(1, W - bw))
                bh = rint(hlo, hhi)
                top = max(0, int(horizon[bx : bx + bw].min()) - bh)
                lab[i, top : int(horizon[bx : bx + bw].max()), bx : bx + bw] = cls

        for c, col in base.items():
            m = lab[i] == c
            for ch in range(3):
                rgb[i, ch][m] = col[ch]
        win = (torch.arange(H).view(-1, 1) % 3 == 0) & (torch.arange(W).view(1, -1) % 3 == 0)
        rgb[i][:, (lab[i] == BUILDING) & win] -= 45  # windows
        foliage = lab[i] == TREE
        if bool(foliage.any()):
            rgb[i][:, foliage] += torch.randn(3, int(foliage.sum()), generator=g) * 28
        lane = (
            (lab[i] == ROAD)
            & (torch.arange(W).view(1, -1) % 7 < 2)
            & (torch.arange(H).view(-1, 1) > H * 0.75)
        )
        rgb[i][:, lane] += 40  # lane markings
        rgb[i] += torch.randn(3, H, W, generator=g) * 10
        rgb[i] += float(torch.randn(1, generator=g)) * 18  # illumination
    rgb = rgb.clamp(0, 255) / 255.0
    dev = torch.device(device) if device else None
    return (rgb.to(dev), lab.to(dev)) if dev else (rgb, lab)

balanced_class_probabilities

balanced_class_probabilities(labels: Tensor, n_classes: Optional[int] = None, ignore_index: Optional[int] = None, floor: float = 0.0) -> Tensor

Per-class feedback probabilities that equalise how often each class is learned from.

Dense targets are dominated by whichever class covers the most pixels, and a Tsetlin machine has no loss to reweight — the only lever is which pixels produce feedback. This returns p_k = min_j freq_j / freq_k (so the rarest class keeps probability 1), ready to pass as class_feedback_p to :class:~torchtsetlin.models.SegmentationTsetlinMachine.

Parameters:

Name	Type	Description	Default
labels	Tensor	any tensor of integer labels (a label map or a flat vector).	required
n_classes	Optional[int]	number of classes (inferred from the maximum label if omitted).	None
ignore_index	Optional[int]	label to exclude from the frequency count.	None
floor	float	lower bound on the returned probabilities; raising it (e.g. 0.05) keeps a very frequent class from being starved of feedback altogether.	0.0

Source code in src/torchtsetlin/data/datasets.py

def balanced_class_probabilities(
    labels: Tensor,
    n_classes: Optional[int] = None,
    ignore_index: Optional[int] = None,
    floor: float = 0.0,
) -> Tensor:
    """Per-class feedback probabilities that equalise how often each class is learned from.

    Dense targets are dominated by whichever class covers the most pixels, and a Tsetlin
    machine has no loss to reweight — the only lever is *which pixels produce feedback*.
    This returns ``p_k = min_j freq_j / freq_k`` (so the rarest class keeps probability 1),
    ready to pass as ``class_feedback_p`` to
    :class:`~torchtsetlin.models.SegmentationTsetlinMachine`.

    Args:
        labels: any tensor of integer labels (a label map or a flat vector).
        n_classes: number of classes (inferred from the maximum label if omitted).
        ignore_index: label to exclude from the frequency count.
        floor: lower bound on the returned probabilities; raising it (e.g. ``0.05``) keeps a
            very frequent class from being starved of feedback altogether.
    """
    y = labels.reshape(-1).long()
    if ignore_index is not None:
        y = y[y != int(ignore_index)]
    K = int(n_classes) if n_classes is not None else (int(y.max()) + 1 if y.numel() else 1)
    counts = torch.bincount(y, minlength=K).to(torch.float64)[:K]
    present = counts[counts > 0]
    if present.numel() == 0:
        return torch.ones(K, dtype=torch.float32, device=labels.device)
    p = torch.where(counts > 0, present.min() / counts.clamp_min(1), torch.ones_like(counts))
    return p.clamp(min=float(floor), max=1.0).to(torch.float32).to(labels.device)

load_torchvision_boolean

load_torchvision_boolean(dataset_cls, root: str, train: bool, encoder: Optional[Callable] = None, threshold: float = 0.3, device=None, download: bool = True, max_samples: Optional[int] = None, **kwargs) -> Tuple[Tensor, Tensor]

Load a torchvision dataset fully into memory and booleanize it.

encoder receives a float tensor (N, C, H, W) scaled to [0, 1] and must return a Boolean tensor. By default pixels are thresholded at threshold (0.3, as in the original MNIST experiments). Returns (x_bool, y_long) on device.

Source code in src/torchtsetlin/data/datasets.py

def load_torchvision_boolean(
    dataset_cls,
    root: str,
    train: bool,
    encoder: Optional[Callable] = None,
    threshold: float = 0.3,
    device=None,
    download: bool = True,
    max_samples: Optional[int] = None,
    **kwargs,
) -> Tuple[Tensor, Tensor]:
    """Load a torchvision dataset fully into memory and booleanize it.

    ``encoder`` receives a float tensor ``(N, C, H, W)`` scaled to ``[0, 1]`` and must return
    a Boolean tensor. By default pixels are thresholded at ``threshold`` (0.3, as in the
    original MNIST experiments). Returns ``(x_bool, y_long)`` on ``device``.
    """
    ds = dataset_cls(root=root, train=train, download=download, **kwargs)
    data = ds.data
    if not isinstance(data, Tensor):
        data = torch.as_tensor(np.asarray(data))
    if data.dim() == 3:  # (N, H, W)
        data = data.unsqueeze(1)
    elif data.dim() == 4 and data.shape[-1] in (1, 3):  # (N, H, W, C) -> (N, C, H, W)
        data = data.permute(0, 3, 1, 2)
    targets = torch.as_tensor(np.asarray(ds.targets)).long()
    if max_samples is not None:
        data, targets = data[:max_samples], targets[:max_samples]
    x = data.float() / 255.0
    if device is not None:
        x, targets = x.to(device), targets.to(device)
    xb = encoder(x) if encoder is not None else (x > threshold)
    return xb, targets

load_mnist_boolean

load_mnist_boolean(root: str = './data', train: bool = True, threshold: float = 0.3, encoder: Optional[Callable] = None, device=None, flatten: bool = False, max_samples: Optional[int] = None) -> Tuple[Tensor, Tensor]

MNIST as Boolean tensors (N, 1, 28, 28) (or (N, 784) with flatten=True). Requires torchvision.

Source code in src/torchtsetlin/data/datasets.py

def load_mnist_boolean(
    root: str = "./data",
    train: bool = True,
    threshold: float = 0.3,
    encoder: Optional[Callable] = None,
    device=None,
    flatten: bool = False,
    max_samples: Optional[int] = None,
) -> Tuple[Tensor, Tensor]:
    """MNIST as Boolean tensors ``(N, 1, 28, 28)`` (or ``(N, 784)`` with ``flatten=True``).
    Requires ``torchvision``."""
    try:
        from torchvision.datasets import MNIST
    except ImportError as e:  # pragma: no cover
        raise ImportError("torchvision is required for load_mnist_boolean") from e
    x, y = load_torchvision_boolean(
        MNIST, root, train, encoder=encoder, threshold=threshold, device=device, max_samples=max_samples
    )
    if flatten:
        x = x.reshape(x.shape[0], -1)
    return x, y

load_segmentation_boolean

load_segmentation_boolean(dataset: Dataset, encoder: Optional[Callable] = None, size: Optional[Union[int, Tuple[int, int]]] = None, n_bits: int = 4, device: Optional[Union[str, device]] = None, max_samples: Optional[int] = None, ignore_index: int = 255, remap: Optional[dict] = None) -> Tuple[Tensor, Tensor]

Load a torchvision segmentation dataset into memory as Boolean planes + label maps.

Segmentation datasets yield (image, mask) PIL pairs rather than the packed tensors :func:load_torchvision_boolean expects, and the mask must be resized with nearest interpolation — bilinear would invent class indices that do not exist. This handles both.

Parameters:

Name	Type	Description	Default
dataset	Dataset	an instantiated torchvision segmentation dataset, e.g. VOCSegmentation(root, image_set="train", download=True) or Cityscapes(root, split="train", target_type="semantic").	required
encoder	Optional[Callable]	maps a float image batch (N, 3, H, W) in [0, 1] to Boolean planes. Defaults to a :class:~torchtsetlin.data.ColorThermometerEncoder with n_bits levels per channel, which keeps the ordering between intensities.	None
size	Optional[Union[int, Tuple[int, int]]]	H or (H, W) to resize to. Dense Tsetlin machines evaluate one patch per pixel, so full-resolution frames are expensive — 64-128 px is a sane starting point. None keeps the native size (all images must then already match).	None
n_bits	int	thermometer levels per channel for the default encoder.	4
device	Optional[Union[str, device]]	device for the returned tensors.	None
max_samples	Optional[int]	stop after this many images.	None
ignore_index	int	mask value to keep as the void class. Pass the same value as the model's ignore_index.	255
remap	Optional[dict]	optional {original_label: new_label} mapping applied to the masks, for collapsing a dataset's classes into a coarser set.	None

Returns:

Type	Description
Tuple[Tensor, Tensor]	(x_bool, labels) with x_bool (N, Z, H, W) and labels (N, H, W).

Source code in src/torchtsetlin/data/datasets.py

def load_segmentation_boolean(
    dataset: Dataset,
    encoder: Optional[Callable] = None,
    size: Optional[Union[int, Tuple[int, int]]] = None,
    n_bits: int = 4,
    device: Optional[Union[str, torch.device]] = None,
    max_samples: Optional[int] = None,
    ignore_index: int = 255,
    remap: Optional[dict] = None,
) -> Tuple[Tensor, Tensor]:
    """Load a torchvision *segmentation* dataset into memory as Boolean planes + label maps.

    Segmentation datasets yield ``(image, mask)`` PIL pairs rather than the packed tensors
    :func:`load_torchvision_boolean` expects, and the mask must be resized with **nearest**
    interpolation — bilinear would invent class indices that do not exist. This handles both.

    Args:
        dataset: an instantiated torchvision segmentation dataset, e.g.
            ``VOCSegmentation(root, image_set="train", download=True)`` or
            ``Cityscapes(root, split="train", target_type="semantic")``.
        encoder: maps a float image batch ``(N, 3, H, W)`` in ``[0, 1]`` to Boolean planes.
            Defaults to a :class:`~torchtsetlin.data.ColorThermometerEncoder` with ``n_bits``
            levels per channel, which keeps the ordering between intensities.
        size: ``H`` or ``(H, W)`` to resize to. Dense Tsetlin machines evaluate one patch per
            pixel, so full-resolution frames are expensive — 64-128 px is a sane starting
            point. ``None`` keeps the native size (all images must then already match).
        n_bits: thermometer levels per channel for the default encoder.
        device: device for the returned tensors.
        max_samples: stop after this many images.
        ignore_index: mask value to keep as the void class. Pass the same value as the
            model's ``ignore_index``.
        remap: optional ``{original_label: new_label}`` mapping applied to the masks, for
            collapsing a dataset's classes into a coarser set.

    Returns:
        ``(x_bool, labels)`` with ``x_bool`` ``(N, Z, H, W)`` and ``labels`` ``(N, H, W)``.
    """
    try:
        from torchvision.transforms import functional as VF
    except ImportError as e:  # pragma: no cover
        raise ImportError("torchvision is required for load_segmentation_boolean") from e

    hw = None if size is None else ((size, size) if isinstance(size, int) else tuple(size))
    n = len(dataset) if max_samples is None else min(len(dataset), int(max_samples))
    imgs, masks = [], []
    for i in range(n):
        img, mask = dataset[i][0], dataset[i][1]
        if hw is not None:
            img = VF.resize(img, list(hw))
            mask = VF.resize(mask, list(hw), interpolation=VF.InterpolationMode.NEAREST)
        img_t = img if isinstance(img, Tensor) else VF.pil_to_tensor(img)
        mask_t = mask if isinstance(mask, Tensor) else VF.pil_to_tensor(mask)
        imgs.append(img_t.float() / 255.0)
        masks.append(mask_t.squeeze(0).long())
    x = torch.stack(imgs)
    y = torch.stack(masks)
    if remap:
        out = torch.full_like(y, int(ignore_index))
        for src, dst in remap.items():
            out[y == int(src)] = int(dst)
        y = out
    if device is not None:
        x, y = x.to(device), y.to(device)
    if encoder is None:
        from .encoders import ColorThermometerEncoder

        encoder = ColorThermometerEncoder(n_bits=n_bits, value_range=(0.0, 1.0))
        if device is not None:
            encoder = encoder.to(device)
    return encoder(x), y

load_voc_segmentation_boolean

load_voc_segmentation_boolean(root: str = './data', image_set: str = 'train', size: Union[int, Tuple[int, int]] = 96, n_bits: int = 4, device=None, download: bool = True, max_samples: Optional[int] = None, **kwargs) -> Tuple[Tensor, Tensor]

Pascal VOC 2012 segmentation as Boolean planes + label maps (21 classes + void).

Void pixels keep the value 255, so pass ignore_index=255 to the model. Requires torchvision.

Source code in src/torchtsetlin/data/datasets.py

def load_voc_segmentation_boolean(
    root: str = "./data",
    image_set: str = "train",
    size: Union[int, Tuple[int, int]] = 96,
    n_bits: int = 4,
    device=None,
    download: bool = True,
    max_samples: Optional[int] = None,
    **kwargs,
) -> Tuple[Tensor, Tensor]:
    """Pascal VOC 2012 segmentation as Boolean planes + label maps (21 classes + void).

    Void pixels keep the value ``255``, so pass ``ignore_index=255`` to the model. Requires
    ``torchvision``.
    """
    try:
        from torchvision.datasets import VOCSegmentation
    except ImportError as e:  # pragma: no cover
        raise ImportError("torchvision is required for load_voc_segmentation_boolean") from e
    ds = VOCSegmentation(root=root, image_set=image_set, download=download, **kwargs)
    return load_segmentation_boolean(
        ds, size=size, n_bits=n_bits, device=device, max_samples=max_samples, ignore_index=255
    )