From 0fe1b37b40924087d3c0458245f137a05e0f2787 Mon Sep 17 00:00:00 2001 From: Aaron Zolnai-Lucas Date: Thu, 22 Oct 2026 04:23:24 +0100 Subject: [PATCH 1/6] guide for creating a tabicl-style tfm --- docs/source/index.md | 1 + docs/source/nanotabicl.md | 381 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 382 insertions(+) create mode 100644 docs/source/nanotabicl.md diff --git a/docs/source/index.md b/docs/source/index.md index feb4079a09..647ea2499d 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -32,6 +32,7 @@ tensor icl processing nim +nanotabicl ``` ```{toctree} diff --git a/docs/source/nanotabicl.md b/docs/source/nanotabicl.md new file mode 100644 index 0000000000..f26cfc7d7c --- /dev/null +++ b/docs/source/nanotabicl.md @@ -0,0 +1,381 @@ +# Building a Tabular Foundation Model + +This guide builds a minimal TabICL-inspired tabular foundation model, `NanoTabICL`, by composing building blocks from the `sdm` library. +Like other in-context learning models, `NanoTabICL` receives labeled *context* rows together with unlabeled *query* rows, and predicts the query targets in a single forward pass without updating its weights (see the [In-Context Learning](icl) guide). +The architecture presented here should not be used for training, since several simplifications have been made to present the core ideas - some of them are discussed in the [Extensions](extensions) section. +We drop the leading batch dimension in all tensor shapes for simplicity. +In practice, tabular foundation models are trained in batches. + +The model is a stack of three stages: + +1. An encoder that turns every table row into a fixed-size embedding (a column-wise, then a row-wise transformer). +2. An in-context learning (ICL) transformer that lets query rows learn from labeled context rows. +3. A prediction head. + +`NanoTabICL` is modeled after {py:class}`~sdm.models.TabICLv2` ([Qu et al., 2026](https://arxiv.org/abs/2602.11139)). + +## The MLP Helper + +The following simple feed-forward multilayer perceptron (MLP) helper function keeps the layer definitions brief in the sections below: + +```python +from typing import Any + +import sdm +import sdm.nn + +import torch +from torch import nn +from torch import Tensor + + +def mlp(channels: int, out_channels: int | None = None) -> nn.Sequential: + return nn.Sequential( + nn.LayerNorm(channels), + nn.Linear(channels, 2 * channels), + nn.GELU(), + nn.Linear(2 * channels, out_channels if out_channels is not None else channels) + ) +``` + +## Encoder + +The encoder embeds each table row, and consists of a column-wise transformer followed by a row-wise transformer. +Its output can be interpreted as one embedding per row. + +### Column Transformer + +The first encoder stage is a column-wise set transformer, which operates on one embedding per table cell. + +Set transformers ([Lee et al., 2018](https://arxiv.org/abs/1810.00825)) are order-invariant transformers, which is desirable because predictions should not depend on row permutations. +A naive set transformer would remove any sequential masking and positional encoding from the classic transformer architecture. +Since self-attention grows quadratically with set size, however, full self-attention across all cells of a column would not be scalable. + +The key innovation of set transformers is to learn a low, fixed number of inducing points. +These points first attend to the set elements, before the set elements attend to the induced output. +The inducing points themselves are learnable parameters. +A set transformer can be seen as a compressed version of full self-attention. + +Induced self-attention blocks are made up of two sequential transformer blocks: + +1. The inducing block: the inducing points are learned by attending to the input set, in this case the column values. +2. The output block: the input set (the column values) acts as the query and attends to the inducing points' output from step 1, so every cell reads a compressed summary of its whole column instead of attending to every other cell. + +The column stage is the first place where information flows between rows. +We therefore restrict the keys and values of the inducing block to the context rows, while the output block still runs for every row, context or query. +This prevents queries within a batch from influencing each other, and allows caching across query rows (see [KV Caching](kv-caching)). + +{py:class}`~sdm.nn.InducedTransformerBlock` provides a highly reusable and configurable block for induced self-attention, built from two {py:class}`~sdm.nn.TransformerBlock` modules. + +Tensor dimension key: + +- `R`: the number of rows. +- `C`: the number of columns. +- `d`: the embedding dimension. + +```python +class ColumnTransformer(nn.Module): + def __init__( + self, + channels: int, + num_inducing_points: int, + num_heads: int, + num_layers: int, + inducing_block_kw: dict[str, Any] | None = None, + output_block_kw: dict[str, Any] | None = None, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList([ + sdm.nn.InducedTransformerBlock( + channels=channels, + num_inducing_points=num_inducing_points, + inducing_block=sdm.nn.TransformerBlock( + channels, num_heads, mlp(channels), **(inducing_block_kw or {}) + ), + output_block=sdm.nn.TransformerBlock( + channels, num_heads, mlp(channels), **(output_block_kw or {}) + ), + ) for _ in range(num_layers) + ]) + + def forward(self, x: Tensor, n_ctx: int) -> Tensor: + # x: (R, C, d), where the first n_ctx rows are context rows and the rest are query rows + x = x.transpose(0, 1) # (C, R, d): since transformer block attends over second last dimension + for block in self.blocks: + x = block(x, key_value=x[:, :n_ctx]) # inducing points only attend to context rows + return x.transpose(0, 1) # (R, C, d) +``` + +The learned inducing vectors of each block are available through its `inducing_points` property. + +### Row Transformer + +Since TabICL uses column-wise, then row-wise attention, we next perform self-attention across the column vectors of each row. +As the number of columns in a dataset is usually much smaller than the number of rows, we can use a full self-attention block ({py:class}`~sdm.nn.TransformerBlock`) here. + +The number of column and row transformer blocks can be repeated. +In TabICLv2, 3 column attention layers are stacked, followed by 3 row attention layers. +In {py:class}`~sdm.models.KumoTabular`, the column and row attention blocks are interleaved. + +The final row transformer block is the embedding layer before the ICL module. +Hence, it must output a fixed-size tensor that is independent of the number of table columns, in order to decouple the ICL stage from the table size. +To achieve this, we aggregate each row into a fixed number of learned "readout" vectors (also called CLS tokens, borrowing from the NLP literature). +This is similar to the first half of a set transformer, in that we compress to a fixed set of vectors. +However, here we have full self-attention at each layer, and only read out the CLS tokens of the last layer. +We default to 4 readout tokens, as TabICL does. + +Together, the column-then-row attention blocks form the encoder of our tabular foundation model. + +The readout tokens are initialized from a truncated normal distribution, to break the symmetry between readout tokens. +Ultimately, they are concatenated into a single, large vector for the ICL stage. +Using multiple readout tokens allows more information to pass through this bottleneck. + +Tensor dimension key: + +- `R`: the number of rows. +- `C`: the number of columns. +- `M`: the number of readout tokens. +- `d`: the embedding dimension. + +```python +class RowTransformer(nn.Module): + def __init__( + self, + channels: int, + num_layers: int, + num_heads: int, + num_readout_tokens: int = 4, + **transformer_block_kw, + ) -> None: + super().__init__() + self.readout_tokens = nn.Parameter(torch.empty(num_readout_tokens, channels)) + nn.init.trunc_normal_(self.readout_tokens, std=0.02) + + self.blocks = nn.ModuleList([ + sdm.nn.TransformerBlock(channels, num_heads, mlp(channels), **transformer_block_kw) + for _ in range(num_layers) + ]) + + @property + def num_readout_tokens(self) -> int: + return self.readout_tokens.size(0) + + def forward(self, x: Tensor) -> Tensor: + # x: (R, C, d) + readout = self.readout_tokens.expand(x.size(0), -1, -1) # (R, M, d) + x = torch.cat([readout, x], dim=1) # (R, M + C, d) + for block in self.blocks: + x = block(x) + return x[:, :self.num_readout_tokens, :] # (R, M, d) +``` + +## In-Context Learning Block + +Now that we have embeddings for each context and query row from the encoder, we add an in-context learning stage. +This is again a sequence of {py:class}`~sdm.nn.TransformerBlock` modules, with one key mechanism: the attention mask, which follows two rules: + +1. Query rows may only attend to context rows. +2. Context rows may attend to each other, but not to query rows. + +Rule 1 ensures that predictions are stable across batches: they do not depend on which other queries were batched together, and information cannot leak between queries. +Rule 2 ensures that the context computation is independent of the queries, and so can be cached across a large number of query rows (see [KV Caching](kv-caching)). + +Note that, unlike in the encoder stages, we pay the computational price of full self-attention over the context rows here. +This scales quadratically in the number of rows, but in practice it proves its worth empirically. +Each row is represented by a single token, so there is no additional complexity factor from the number of columns, unlike in the encoder. + +Tensor dimension key: + +- `R_ctx`: the number of context rows. +- `R_query`: the number of query rows. +- `M`: the number of readout tokens. +- `d`: the embedding dimension. + +```python +class ICLTransformer(nn.Module): + def __init__( + self, + channels: int, + num_heads: int, + num_layers: int, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList([ + sdm.nn.TransformerBlock(channels, num_heads, mlp(channels)) + for _ in range(num_layers) + ]) + + def forward(self, x_context: Tensor, x_query: Tensor) -> Tensor: + # x_context: (R_ctx, M * d) label-aware context row embeddings + # x_query: (R_query, M * d) + n_ctx = x_context.size(0) + + x = torch.cat([x_context, x_query], dim=0) # (R_ctx + R_query, M * d) + + # attn_mask = torch.zeros(x.size(0), dtype=torch.bool, device=x.device) + # attn_mask[:n_ctx] = True + + for block in self.blocks: + x = block(query=x, key_value=x[:n_ctx]) + + return x[n_ctx:] # (R_query, M * d) +``` + +## Prediction Head + +After the ICL module, the model needs a way to make predictions. +The head is just an MLP from the ICL output to `out_channels`, so it is not specific to classification: setting `out_channels` to the number of classes gives class logits. +This is the `mlp` helper from above, whose leading layer norm also serves as the final normalization before the head. +For classification, a softmax over the logits gives the class probabilities. + +## Putting It Together + +We can now compose the blocks into a tabular foundation model for classification, made up of three stages in sequence: + +1. The encoder: each cell is first embedded with a shared linear layer, which assumes standardized numerical features, followed by a column-wise, then a row-wise transformer. The output is flattened so that the readout tokens form one larger vector representation of each row, and the embedding scale is normalized to stabilize training in the ICL layers. In TabICL and other architectures, each transformer block also normalizes its attention inputs, but we omit this for simplicity. +2. The ICL module: a transformer stack in which query rows learn from context rows. +3. The prediction head: a classification head that maps the ICL output to class logits. + +Labels enter the model in two places, both added to the context rows only: once to the cell embeddings before the column stage, and once to the flattened row embeddings before the ICL stage. +This is how the model learns from the context: context rows carry their targets as label embeddings, while query rows do not, so the model has to infer the query targets from the context. + +Tensor dimension key: + +- `R`: the number of rows. +- `C`: the number of columns. +- `R_ctx`: the number of context rows. +- `R_query`: the number of query rows. +- `M`: the number of readout tokens. +- `d`: the embedding dimension. + +```python +class NanoTabICL(nn.Module): + """A small, TabICL inspired tabular foundation model.""" + + def __init__( + self, + channels: int, + num_readout_tokens: int = 4, + num_classes: int = 2, + num_heads: int = 4, + num_layers: int = 1, + num_inducing_points: int = 16, + ) -> None: + super().__init__() + icl_channels = num_readout_tokens * channels + + self.cell_embedding = nn.Linear(1, channels) # shared across all cells + self.column_label_embedding = nn.Embedding(num_classes, channels) + self.column_transformer = ColumnTransformer( + channels=channels, + num_inducing_points=num_inducing_points, + num_heads=num_heads, + num_layers=num_layers, + ) + self.row_transformer = RowTransformer( + channels=channels, + num_layers=num_layers, + num_heads=num_heads, + num_readout_tokens=num_readout_tokens, + ) + self.encoder_norm = nn.LayerNorm(channels) + self.icl_label_embedding = nn.Embedding(num_classes, icl_channels) + self.icl_transformer = ICLTransformer( + channels=icl_channels, + num_heads=num_heads, + num_layers=num_layers, + ) + self.prediction_head = mlp(icl_channels, out_channels=num_classes) + + def forward(self, x: Tensor, y_context: Tensor) -> Tensor: + # x: (R, C) numerical features, y_context: (R_ctx,) integer class labels of the context rows + n_ctx = y_context.size(0) + x = self.cell_embedding(x.unsqueeze(-1)) # (R, C, d) + + # Label-aware cell embeddings: context rows only + x_context = x[:n_ctx] + self.column_label_embedding(y_context).unsqueeze(1) # (R_ctx, C, d) + x = torch.cat([x_context, x[n_ctx:]], dim=0) # (R, C, d) + + x = self.column_transformer(x, n_ctx) + x = self.encoder_norm(self.row_transformer(x)).flatten(start_dim=1) # (R, M * d) + + # Label-aware row embeddings: context rows only + x_context = x[:n_ctx] + self.icl_label_embedding(y_context) # (R_ctx, M * d) + x_query = self.icl_transformer(x_context, x[n_ctx:]) # (R_query, M * d) + return self.prediction_head(x_query) # (R_query, num_classes) +``` + +The defaults use a single layer of each kind, whereas TabICLv2 stacks 3 column layers, 3 row layers and 12 ICL layers. +The number of heads must divide both `channels` and `num_readout_tokens * channels`. + +(extensions)= + +## Extensions + +The following are some extra mechanisms included in TabICLv2 which were not covered above. + +### Feature Grouping + +An improvement that TabICLv2 introduces is feature grouping as the first step of the encoder. + +Since columns are embedded first in TabICL, columns with similar distributions can collapse into the same embedding, which makes it hard for the model to distinguish them. +TabICLv2 therefore linearly combines groups of columns, so that pairs of columns are never in the same group. + +A simple grouping that guarantees no overlapping pairs between groups is the following: + +- For each column _i_, group it with columns _i+1_ and _i+3_, wrapping around indices if needed. +- No two groups share 2 columns, due to the unique differences between offsets. + +In general, using shift offsets of `(0, 1, 3, ..., 2**n - 1)`, which form groups of `n + 1` columns, guarantees no overlap when there are at least `2**(n+1) - 1` columns. +In practice, the TabICL authors found no major performance benefit from groups larger than 3, so we'll default to `(0, 1, 3)`. + +After grouping, we apply a shared linear transformation that maps all groups to a fixed embedding dimension, ready for the next layer. +This generalizes the cell embedding of `NanoTabICL` from single columns to groups of columns, so using it amounts to replacing `self.cell_embedding(x.unsqueeze(-1))` in `NanoTabICL` with a `FeatureGrouping(embedding_dim=channels)` applied to `x`. + +Tensor dimension key: + +- `R`: the number of rows. +- `C`: the number of columns. +- `d`: the embedding dimension. + +```python +class FeatureGrouping(nn.Module): + """A simplified feature grouping mechanism for single tables, inspired by TabICLv2.""" + + def __init__( + self, + shift_pattern: tuple[int, ...] = (0, 1, 3), + embedding_dim: int = 64 + ) -> None: + super().__init__() + self.shift_pattern = shift_pattern + self.lin = nn.Linear(len(shift_pattern), embedding_dim) + + def forward(self, x: Tensor) -> Tensor: + assert len(x.shape) == 2 # (R, C) + + # accumulate into (R, C, #offsets) + groups = torch.stack( + [x.roll(offset, dims=1) for offset in self.shift_pattern], + dim=-1, + ) + + return self.lin(groups) # (R, C, d) +``` + +### QASSMax + +Softmax attention spreads out as the number of keys grows: each key gets a smaller weight, and attention can no longer focus on a few relevant rows. +This matters for tabular in-context learning, where the context can range from a handful to many thousands of rows. + +Query-Aware Scalable SoftMax ({py:class}`~sdm.nn.QASSMax`), introduced in TabICLv2 ([Qu et al., 2026](https://arxiv.org/abs/2602.11139)), overcomes this by rescaling the query depending on the number of context rows and on the query itself. + +(kv-caching)= + +### KV Caching + +The same context is often used to predict many batches of query rows, and since the context is usually much larger than a query batch, recomputing it for every batch dominates the cost of prediction. + +The attention masks in our model ensure that context rows never attend to query rows. +The keys and values that the context produces in each layer can therefore be computed once and reused for every batch of query rows - this is what {py:class}`~sdm.cache.Cache` does. +ICL models use it under the hood: {py:meth}`~sdm.models.ICLModel.fit` caches the context once, and each {py:meth}`~sdm.models.ICLModel.predict` call reuses it. From bdc64377c4e276ddc3409853c3ca41cf0abaf93c Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Thu, 22 Oct 2026 11:10:48 +0000 Subject: [PATCH 2/6] Add behavioral tests for the NanoTabICL guide examples --- test/docs/test_nanotabicl.py | 354 +++++++++++++++++++++++++++++++++++ 1 file changed, 354 insertions(+) create mode 100644 test/docs/test_nanotabicl.py diff --git a/test/docs/test_nanotabicl.py b/test/docs/test_nanotabicl.py new file mode 100644 index 0000000000..3b4bc18d8e --- /dev/null +++ b/test/docs/test_nanotabicl.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import re +from pathlib import Path +from typing import Any + +import pytest +import torch +from torch import nn + +from sdm.nn import Attention +from sdm.testing import withCUDA + + +@pytest.fixture(scope="module") +def guide() -> dict[str, Any]: + # Execute the published examples in order so tests cannot drift from them. + path = Path(__file__).parents[2] / "docs/source/nanotabicl.md" + namespace: dict[str, Any] = {} + text = path.read_text() + for match in re.finditer(r"^```python\n(.*?)^```", text, re.M | re.S): + # Preserve guide line numbers in syntax errors and tracebacks. + source = "\n" * text.count("\n", 0, match.start(1)) + match[1] + exec(compile(source, str(path), "exec"), namespace) + return namespace + + +def _initialize_attention(module: nn.Module) -> None: + # SDM zero-initializes attention outputs. Exercise attention rather than + # letting leakage and permutation tests pass through residuals alone. + for child in module.modules(): + if isinstance(child, Attention): + nn.init.eye_(child.out_lin.weight) + module.eval() + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize("out_channels", [None, 1, 3]) +def test_mlp( + guide: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + out_channels: int | None, +) -> None: + module = guide["mlp"](8, out_channels=out_channels).to( + device=device, dtype=dtype + ) + x = torch.randn(2, 3, 8, device=device, dtype=dtype) + + out = module(x) + + assert out.shape == (2, 3, 8 if out_channels is None else out_channels) + assert out.dtype == dtype + assert out.device == device + assert out.isfinite().all() + + +@pytest.mark.parametrize("num_layers", [1, 2]) +@torch.no_grad() +def test_column_transformer_context_only_attention( + guide: dict[str, Any], num_layers: int +) -> None: + module = guide["ColumnTransformer"]( + channels=8, + num_inducing_points=3, + num_heads=2, + num_layers=num_layers, + ) + _initialize_attention(module) + x = torch.randn(7, 3, 8) + n_ctx = 4 + + out = module(x, n_ctx) + assert out.shape == x.shape + + changed_query = x.clone() + changed_query[-1] += torch.randn_like(changed_query[-1]) * 5 + changed_out = module(changed_query, n_ctx) + torch.testing.assert_close(changed_out[:-1], out[:-1]) + assert not torch.allclose(changed_out[-1], out[-1]) + + # Neither the context nor a query may depend on companion queries. + torch.testing.assert_close(module(x[:n_ctx], n_ctx), out[:n_ctx]) + for index in range(n_ctx, x.size(0)): + single = torch.cat([x[:n_ctx], x[index : index + 1]]) + torch.testing.assert_close(module(single, n_ctx)[-1], out[index]) + + changed_context = x.clone() + changed_context[:n_ctx] += torch.randn_like(x[:n_ctx]) * 5 + assert not torch.allclose( + module(changed_context, n_ctx)[n_ctx:], out[n_ctx:] + ) + + +@torch.no_grad() +def test_column_transformer_permutations(guide: dict[str, Any]) -> None: + module = guide["ColumnTransformer"]( + channels=8, num_inducing_points=3, num_heads=2, num_layers=2 + ) + _initialize_attention(module) + x = torch.randn(7, 3, 8) + rows = torch.tensor([2, 0, 3, 1, 6, 4, 5]) + columns = torch.tensor([2, 0, 1]) + + out = module(x, n_ctx=4) + + torch.testing.assert_close(module(x[rows], n_ctx=4), out[rows]) + torch.testing.assert_close(module(x[:, columns], n_ctx=4), out[:, columns]) + + +@pytest.mark.parametrize("num_readout_tokens", [1, 4]) +@torch.no_grad() +def test_row_transformer_readout( + guide: dict[str, Any], num_readout_tokens: int +) -> None: + module = guide["RowTransformer"]( + channels=8, + num_layers=2, + num_heads=2, + num_readout_tokens=num_readout_tokens, + ) + _initialize_attention(module) + assert module.num_readout_tokens == num_readout_tokens + + # The same encoder produces a fixed readout width for different tables. + for num_columns in (1, 5): + x = torch.randn(3, num_columns, 8) + out = module(x) + + assert out.shape == (3, num_readout_tokens, 8) + torch.testing.assert_close(module(x.flip(1)), out) + torch.testing.assert_close(module(x.flip(0)), out.flip(0)) + torch.testing.assert_close(module(x[:1]), out[:1]) + + changed = x.clone() + changed[0] += torch.randn_like(changed[0]) * 5 + changed_out = module(changed) + torch.testing.assert_close(changed_out[1:], out[1:]) + assert not torch.allclose(changed_out[0], out[0]) + + +@pytest.mark.parametrize("num_layers", [1, 2]) +@torch.no_grad() +def test_icl_transformer_context_and_query_isolation( + guide: dict[str, Any], num_layers: int +) -> None: + module = guide["ICLTransformer"]( + channels=8, num_heads=2, num_layers=num_layers + ) + _initialize_attention(module) + context = torch.randn(4, 8) + query = torch.randn(3, 8) + out = module(context, query) + + assert out.shape == query.shape + for index in range(query.size(0)): + torch.testing.assert_close( + module(context, query[index : index + 1]), out[index : index + 1] + ) + torch.testing.assert_close(module(context.flip(0), query), out) + torch.testing.assert_close(module(context, query.flip(0)), out.flip(0)) + + changed_query = query.clone() + changed_query[-1] += torch.randn_like(query[-1]) * 5 + changed_out = module(context, changed_query) + torch.testing.assert_close(changed_out[:-1], out[:-1]) + assert not torch.allclose(changed_out[-1], out[-1]) + assert not torch.allclose( + module(context + torch.randn_like(context), query), out + ) + assert module(context, query[:0]).shape == (0, 8) + + +@pytest.fixture +def model(guide: dict[str, Any]) -> nn.Module: + module = guide["NanoTabICL"]( + channels=8, + num_readout_tokens=2, + num_classes=3, + num_heads=2, + num_layers=2, + num_inducing_points=3, + ) + _initialize_attention(module) + return module + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize( + ("num_context", "num_query", "num_columns", "num_classes"), + [(1, 1, 1, 2), (4, 3, 5, 3), (4, 0, 2, 3)], +) +@torch.no_grad() +def test_nanotabicl_forward( + guide: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + num_context: int, + num_query: int, + num_columns: int, + num_classes: int, +) -> None: + module = guide["NanoTabICL"](channels=8, num_classes=num_classes).to( + device=device, dtype=dtype + ) + _initialize_attention(module) + x = torch.randn( + num_context + num_query, num_columns, device=device, dtype=dtype + ) + y = torch.arange(num_context, device=device) % num_classes + + out = module(x, y) + + assert out.shape == (num_query, num_classes) + assert out.dtype == dtype + assert out.device == device + assert out.isfinite().all() + probabilities = out.softmax(dim=-1) + torch.testing.assert_close( + probabilities.sum(dim=-1), + torch.ones(num_query, device=device, dtype=dtype), + ) + + +@torch.no_grad() +def test_nanotabicl_permutations(model: nn.Module) -> None: + x = torch.randn(7, 5) + y = torch.tensor([0, 1, 2, 1]) + context_order = torch.tensor([2, 0, 3, 1]) + query_order = torch.tensor([2, 0, 1]) + out = model(x, y) + + permuted_context = torch.cat([x[:4][context_order], x[4:]]) + torch.testing.assert_close(model(permuted_context, y[context_order]), out) + permuted_query = torch.cat([x[:4], x[4:][query_order]]) + torch.testing.assert_close(model(permuted_query, y), out[query_order]) + torch.testing.assert_close(model(x.flip(1), y), out) + + +@torch.no_grad() +def test_nanotabicl_query_batch_independence(model: nn.Module) -> None: + x = torch.randn(7, 5) + y = torch.tensor([0, 1, 2, 1]) + out = model(x, y) + + separate = torch.cat( + [model(torch.cat([x[:4], row.unsqueeze(0)]), y) for row in x[4:]] + ) + torch.testing.assert_close(separate, out) + + changed = x.clone() + changed[-1] += torch.randn_like(changed[-1]) * 5 + changed_out = model(changed, y) + torch.testing.assert_close(changed_out[:-1], out[:-1]) + assert not torch.allclose(changed_out[-1], out[-1]) + + +@pytest.mark.parametrize("label_stage", ["column", "icl"]) +@torch.no_grad() +def test_nanotabicl_uses_context_labels( + model: nn.Module, label_stage: str +) -> None: + # Isolate each documented label injection so dropping either is caught. + if label_stage == "column": + model.icl_label_embedding.weight.zero_() + else: + model.column_label_embedding.weight.zero_() + x = torch.randn(7, 5) + y = torch.zeros(4, dtype=torch.long) + + assert not torch.allclose(model(x, y), model(x, y + 1)) + + +@pytest.mark.parametrize("label", [-1, 3]) +def test_nanotabicl_invalid_context_label( + model: nn.Module, label: int +) -> None: + with pytest.raises(IndexError): + model(torch.randn(3, 2), torch.tensor([0, label])) + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize( + ("num_columns", "shift_pattern"), + [ + (7, (0, 1, 3)), + (2, (0, 1, 3)), + (1, (0, 1, 3)), + (4, (0,)), + (4, (-1, 0, 6)), + ], +) +@torch.no_grad() +def test_feature_grouping_offsets( + guide: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + num_columns: int, + shift_pattern: tuple[int, ...], +) -> None: + module = guide["FeatureGrouping"]( + shift_pattern=shift_pattern, embedding_dim=len(shift_pattern) + ).to(device=device, dtype=dtype) + module.lin.weight.copy_( + torch.eye(len(shift_pattern), device=device, dtype=dtype) + ) + module.lin.bias.zero_() + x = torch.arange(2 * num_columns, device=device, dtype=dtype).reshape( + 2, num_columns + ) + + out = module(x) + + assert out.shape == (2, num_columns, len(shift_pattern)) + assert out.dtype == dtype + assert out.device == device + # An independent indexing oracle checks shift direction and wraparound. + for column in range(num_columns): + expected = x[ + :, [(column - shift) % num_columns for shift in shift_pattern] + ] + torch.testing.assert_close(out[:, column], expected) + + +@torch.no_grad() +def test_feature_grouping_projection(guide: dict[str, Any]) -> None: + module = guide["FeatureGrouping"](embedding_dim=2) + module.lin.weight.copy_(torch.tensor([[1.0, 2.0, 3.0], [-1.0, 0.0, 1.0]])) + module.lin.bias.copy_(torch.tensor([0.5, -2.0])) + x = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]) + + out = module(x) + + expected = torch.tensor( + [ + [[15.5, -1.0], [13.5, -1.0], [19.5, -1.0], [13.5, -5.0]], + [[39.5, -1.0], [37.5, -1.0], [43.5, -1.0], [37.5, -5.0]], + ] + ) + torch.testing.assert_close(out, expected) + + +@pytest.mark.parametrize("shape", [(3,), (2, 3, 4)]) +def test_feature_grouping_rejects_non_matrix( + guide: dict[str, Any], shape: tuple[int, ...] +) -> None: + module = guide["FeatureGrouping"]() + + with pytest.raises(AssertionError): + module(torch.zeros(shape)) From 15bbedff66c4cb33dd6fdeb91871a016d2ea3595 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:50:45 -0700 Subject: [PATCH 3/6] Narrow label embeddings in NanoTabICL context-label test --- test/docs/test_nanotabicl.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/docs/test_nanotabicl.py b/test/docs/test_nanotabicl.py index 3b4bc18d8e..c3a2482a2c 100644 --- a/test/docs/test_nanotabicl.py +++ b/test/docs/test_nanotabicl.py @@ -265,9 +265,11 @@ def test_nanotabicl_uses_context_labels( ) -> None: # Isolate each documented label injection so dropping either is caught. if label_stage == "column": - model.icl_label_embedding.weight.zero_() + embedding = model.icl_label_embedding else: - model.column_label_embedding.weight.zero_() + embedding = model.column_label_embedding + assert isinstance(embedding, nn.Embedding) + embedding.weight.zero_() x = torch.randn(7, 5) y = torch.zeros(4, dtype=torch.long) From 49fa45f5f6c437f1ace7131df3134db2f1070305 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Sat, 24 Oct 2026 09:41:19 +0000 Subject: [PATCH 4/6] update --- docs/source/index.md | 1 - docs/source/nanotabicl.md | 7 +- test/docs/test_nanotabicl.py | 347 +---------------------------------- 3 files changed, 11 insertions(+), 344 deletions(-) diff --git a/docs/source/index.md b/docs/source/index.md index f98ba1a41b..9c4a04c214 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -33,7 +33,6 @@ icl processing nanotabicl nim -nanotabicl ``` ```{toctree} diff --git a/docs/source/nanotabicl.md b/docs/source/nanotabicl.md index 0824bd8781..823fca491d 100644 --- a/docs/source/nanotabicl.md +++ b/docs/source/nanotabicl.md @@ -268,9 +268,10 @@ The remaining rows are treated as query rows whose labels the model predicts. Categorical encoding, missing-value handling, and feature scaling are typically performed in a pre-processing stage before the model's {py:meth}`~torch.nn.Module.forward` call. ```python -x = torch.randn(128, 32, device="cuda") -y = torch.randint(0, 10, (64, ), device="cuda") +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +x = torch.randn(128, 32, device=device) +y = torch.randint(0, 10, (64, ), device=device) -model = NanoTabICL().cuda() +model = NanoTabICL().to(device) out = model(x, y) ``` diff --git a/test/docs/test_nanotabicl.py b/test/docs/test_nanotabicl.py index c3a2482a2c..6a7ef3eeb2 100644 --- a/test/docs/test_nanotabicl.py +++ b/test/docs/test_nanotabicl.py @@ -5,352 +5,19 @@ from pathlib import Path from typing import Any -import pytest -import torch -from torch import nn +from torch import Tensor -from sdm.nn import Attention -from sdm.testing import withCUDA - -@pytest.fixture(scope="module") -def guide() -> dict[str, Any]: - # Execute the published examples in order so tests cannot drift from them. +def test_nanotabicl() -> None: path = Path(__file__).parents[2] / "docs/source/nanotabicl.md" - namespace: dict[str, Any] = {} text = path.read_text() + namespace: dict[str, Any] = {} + for match in re.finditer(r"^```python\n(.*?)^```", text, re.M | re.S): - # Preserve guide line numbers in syntax errors and tracebacks. source = "\n" * text.count("\n", 0, match.start(1)) + match[1] exec(compile(source, str(path), "exec"), namespace) - return namespace - - -def _initialize_attention(module: nn.Module) -> None: - # SDM zero-initializes attention outputs. Exercise attention rather than - # letting leakage and permutation tests pass through residuals alone. - for child in module.modules(): - if isinstance(child, Attention): - nn.init.eye_(child.out_lin.weight) - module.eval() - - -@withCUDA -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize("out_channels", [None, 1, 3]) -def test_mlp( - guide: dict[str, Any], - device: torch.device, - dtype: torch.dtype, - out_channels: int | None, -) -> None: - module = guide["mlp"](8, out_channels=out_channels).to( - device=device, dtype=dtype - ) - x = torch.randn(2, 3, 8, device=device, dtype=dtype) - - out = module(x) - - assert out.shape == (2, 3, 8 if out_channels is None else out_channels) - assert out.dtype == dtype - assert out.device == device - assert out.isfinite().all() - - -@pytest.mark.parametrize("num_layers", [1, 2]) -@torch.no_grad() -def test_column_transformer_context_only_attention( - guide: dict[str, Any], num_layers: int -) -> None: - module = guide["ColumnTransformer"]( - channels=8, - num_inducing_points=3, - num_heads=2, - num_layers=num_layers, - ) - _initialize_attention(module) - x = torch.randn(7, 3, 8) - n_ctx = 4 - - out = module(x, n_ctx) - assert out.shape == x.shape - - changed_query = x.clone() - changed_query[-1] += torch.randn_like(changed_query[-1]) * 5 - changed_out = module(changed_query, n_ctx) - torch.testing.assert_close(changed_out[:-1], out[:-1]) - assert not torch.allclose(changed_out[-1], out[-1]) - - # Neither the context nor a query may depend on companion queries. - torch.testing.assert_close(module(x[:n_ctx], n_ctx), out[:n_ctx]) - for index in range(n_ctx, x.size(0)): - single = torch.cat([x[:n_ctx], x[index : index + 1]]) - torch.testing.assert_close(module(single, n_ctx)[-1], out[index]) - - changed_context = x.clone() - changed_context[:n_ctx] += torch.randn_like(x[:n_ctx]) * 5 - assert not torch.allclose( - module(changed_context, n_ctx)[n_ctx:], out[n_ctx:] - ) - - -@torch.no_grad() -def test_column_transformer_permutations(guide: dict[str, Any]) -> None: - module = guide["ColumnTransformer"]( - channels=8, num_inducing_points=3, num_heads=2, num_layers=2 - ) - _initialize_attention(module) - x = torch.randn(7, 3, 8) - rows = torch.tensor([2, 0, 3, 1, 6, 4, 5]) - columns = torch.tensor([2, 0, 1]) - - out = module(x, n_ctx=4) - - torch.testing.assert_close(module(x[rows], n_ctx=4), out[rows]) - torch.testing.assert_close(module(x[:, columns], n_ctx=4), out[:, columns]) - - -@pytest.mark.parametrize("num_readout_tokens", [1, 4]) -@torch.no_grad() -def test_row_transformer_readout( - guide: dict[str, Any], num_readout_tokens: int -) -> None: - module = guide["RowTransformer"]( - channels=8, - num_layers=2, - num_heads=2, - num_readout_tokens=num_readout_tokens, - ) - _initialize_attention(module) - assert module.num_readout_tokens == num_readout_tokens - - # The same encoder produces a fixed readout width for different tables. - for num_columns in (1, 5): - x = torch.randn(3, num_columns, 8) - out = module(x) - - assert out.shape == (3, num_readout_tokens, 8) - torch.testing.assert_close(module(x.flip(1)), out) - torch.testing.assert_close(module(x.flip(0)), out.flip(0)) - torch.testing.assert_close(module(x[:1]), out[:1]) - - changed = x.clone() - changed[0] += torch.randn_like(changed[0]) * 5 - changed_out = module(changed) - torch.testing.assert_close(changed_out[1:], out[1:]) - assert not torch.allclose(changed_out[0], out[0]) - - -@pytest.mark.parametrize("num_layers", [1, 2]) -@torch.no_grad() -def test_icl_transformer_context_and_query_isolation( - guide: dict[str, Any], num_layers: int -) -> None: - module = guide["ICLTransformer"]( - channels=8, num_heads=2, num_layers=num_layers - ) - _initialize_attention(module) - context = torch.randn(4, 8) - query = torch.randn(3, 8) - out = module(context, query) - - assert out.shape == query.shape - for index in range(query.size(0)): - torch.testing.assert_close( - module(context, query[index : index + 1]), out[index : index + 1] - ) - torch.testing.assert_close(module(context.flip(0), query), out) - torch.testing.assert_close(module(context, query.flip(0)), out.flip(0)) - changed_query = query.clone() - changed_query[-1] += torch.randn_like(query[-1]) * 5 - changed_out = module(context, changed_query) - torch.testing.assert_close(changed_out[:-1], out[:-1]) - assert not torch.allclose(changed_out[-1], out[-1]) - assert not torch.allclose( - module(context + torch.randn_like(context), query), out - ) - assert module(context, query[:0]).shape == (0, 8) - - -@pytest.fixture -def model(guide: dict[str, Any]) -> nn.Module: - module = guide["NanoTabICL"]( - channels=8, - num_readout_tokens=2, - num_classes=3, - num_heads=2, - num_layers=2, - num_inducing_points=3, - ) - _initialize_attention(module) - return module - - -@withCUDA -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize( - ("num_context", "num_query", "num_columns", "num_classes"), - [(1, 1, 1, 2), (4, 3, 5, 3), (4, 0, 2, 3)], -) -@torch.no_grad() -def test_nanotabicl_forward( - guide: dict[str, Any], - device: torch.device, - dtype: torch.dtype, - num_context: int, - num_query: int, - num_columns: int, - num_classes: int, -) -> None: - module = guide["NanoTabICL"](channels=8, num_classes=num_classes).to( - device=device, dtype=dtype - ) - _initialize_attention(module) - x = torch.randn( - num_context + num_query, num_columns, device=device, dtype=dtype - ) - y = torch.arange(num_context, device=device) % num_classes - - out = module(x, y) - - assert out.shape == (num_query, num_classes) - assert out.dtype == dtype - assert out.device == device + out = namespace["out"] + assert isinstance(out, Tensor) + assert out.size() == (64, 10) assert out.isfinite().all() - probabilities = out.softmax(dim=-1) - torch.testing.assert_close( - probabilities.sum(dim=-1), - torch.ones(num_query, device=device, dtype=dtype), - ) - - -@torch.no_grad() -def test_nanotabicl_permutations(model: nn.Module) -> None: - x = torch.randn(7, 5) - y = torch.tensor([0, 1, 2, 1]) - context_order = torch.tensor([2, 0, 3, 1]) - query_order = torch.tensor([2, 0, 1]) - out = model(x, y) - - permuted_context = torch.cat([x[:4][context_order], x[4:]]) - torch.testing.assert_close(model(permuted_context, y[context_order]), out) - permuted_query = torch.cat([x[:4], x[4:][query_order]]) - torch.testing.assert_close(model(permuted_query, y), out[query_order]) - torch.testing.assert_close(model(x.flip(1), y), out) - - -@torch.no_grad() -def test_nanotabicl_query_batch_independence(model: nn.Module) -> None: - x = torch.randn(7, 5) - y = torch.tensor([0, 1, 2, 1]) - out = model(x, y) - - separate = torch.cat( - [model(torch.cat([x[:4], row.unsqueeze(0)]), y) for row in x[4:]] - ) - torch.testing.assert_close(separate, out) - - changed = x.clone() - changed[-1] += torch.randn_like(changed[-1]) * 5 - changed_out = model(changed, y) - torch.testing.assert_close(changed_out[:-1], out[:-1]) - assert not torch.allclose(changed_out[-1], out[-1]) - - -@pytest.mark.parametrize("label_stage", ["column", "icl"]) -@torch.no_grad() -def test_nanotabicl_uses_context_labels( - model: nn.Module, label_stage: str -) -> None: - # Isolate each documented label injection so dropping either is caught. - if label_stage == "column": - embedding = model.icl_label_embedding - else: - embedding = model.column_label_embedding - assert isinstance(embedding, nn.Embedding) - embedding.weight.zero_() - x = torch.randn(7, 5) - y = torch.zeros(4, dtype=torch.long) - - assert not torch.allclose(model(x, y), model(x, y + 1)) - - -@pytest.mark.parametrize("label", [-1, 3]) -def test_nanotabicl_invalid_context_label( - model: nn.Module, label: int -) -> None: - with pytest.raises(IndexError): - model(torch.randn(3, 2), torch.tensor([0, label])) - - -@withCUDA -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -@pytest.mark.parametrize( - ("num_columns", "shift_pattern"), - [ - (7, (0, 1, 3)), - (2, (0, 1, 3)), - (1, (0, 1, 3)), - (4, (0,)), - (4, (-1, 0, 6)), - ], -) -@torch.no_grad() -def test_feature_grouping_offsets( - guide: dict[str, Any], - device: torch.device, - dtype: torch.dtype, - num_columns: int, - shift_pattern: tuple[int, ...], -) -> None: - module = guide["FeatureGrouping"]( - shift_pattern=shift_pattern, embedding_dim=len(shift_pattern) - ).to(device=device, dtype=dtype) - module.lin.weight.copy_( - torch.eye(len(shift_pattern), device=device, dtype=dtype) - ) - module.lin.bias.zero_() - x = torch.arange(2 * num_columns, device=device, dtype=dtype).reshape( - 2, num_columns - ) - - out = module(x) - - assert out.shape == (2, num_columns, len(shift_pattern)) - assert out.dtype == dtype - assert out.device == device - # An independent indexing oracle checks shift direction and wraparound. - for column in range(num_columns): - expected = x[ - :, [(column - shift) % num_columns for shift in shift_pattern] - ] - torch.testing.assert_close(out[:, column], expected) - - -@torch.no_grad() -def test_feature_grouping_projection(guide: dict[str, Any]) -> None: - module = guide["FeatureGrouping"](embedding_dim=2) - module.lin.weight.copy_(torch.tensor([[1.0, 2.0, 3.0], [-1.0, 0.0, 1.0]])) - module.lin.bias.copy_(torch.tensor([0.5, -2.0])) - x = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]) - - out = module(x) - - expected = torch.tensor( - [ - [[15.5, -1.0], [13.5, -1.0], [19.5, -1.0], [13.5, -5.0]], - [[39.5, -1.0], [37.5, -1.0], [43.5, -1.0], [37.5, -5.0]], - ] - ) - torch.testing.assert_close(out, expected) - - -@pytest.mark.parametrize("shape", [(3,), (2, 3, 4)]) -def test_feature_grouping_rejects_non_matrix( - guide: dict[str, Any], shape: tuple[int, ...] -) -> None: - module = guide["FeatureGrouping"]() - - with pytest.raises(AssertionError): - module(torch.zeros(shape)) From f7020b2404f0b0af98bc59dadca3be92cd318aad Mon Sep 17 00:00:00 2001 From: rusty1s Date: Sat, 24 Oct 2026 09:42:19 +0000 Subject: [PATCH 5/6] update --- test/docs/test_nanotabicl.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/test/docs/test_nanotabicl.py b/test/docs/test_nanotabicl.py index 6a7ef3eeb2..619c5709b4 100644 --- a/test/docs/test_nanotabicl.py +++ b/test/docs/test_nanotabicl.py @@ -5,8 +5,6 @@ from pathlib import Path from typing import Any -from torch import Tensor - def test_nanotabicl() -> None: path = Path(__file__).parents[2] / "docs/source/nanotabicl.md" @@ -18,6 +16,5 @@ def test_nanotabicl() -> None: exec(compile(source, str(path), "exec"), namespace) out = namespace["out"] - assert isinstance(out, Tensor) assert out.size() == (64, 10) assert out.isfinite().all() From 8ac32ed24f8a7fae8d376ad9c06a876c61ced5d4 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Sat, 24 Oct 2026 09:46:36 +0000 Subject: [PATCH 6/6] update --- docs/source/nanotabicl.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/nanotabicl.md b/docs/source/nanotabicl.md index 823fca491d..99535ceb2c 100644 --- a/docs/source/nanotabicl.md +++ b/docs/source/nanotabicl.md @@ -1,4 +1,4 @@ -# NanoTabICL in SDM +# NanoTabICL The [`sdm.nn`](api/nn) package provides reusable, GPU-accelerated building blocks for developing foundation models for structured data from scratch. In this guide, we show how to build a minimal [TabICLv2](https://arxiv.org/abs/2602.11139)-inspired Tabular Foundation Model (TFM) with it.