TabM

https://github.com/yandex-research/tabm
Admin
TabM is a PyTorch-based Python package providing a simple and powerful tabular deep learning...
Tokens:84,061
Snippets:516
Trust Score:7.3
License:Apache-2.0
Update:1 week ago
Show doc for...
Tokens:
Raw
### Install tabm and rtdl_num_embeddings

Source: https://github.com/yandex-research/tabm/blob/main/example.ipynb

Install the required packages for the TabM example.

```python
!pip install rtdl_num_embeddings
!pip install tabm
```

--------------------------------

### Example: Create and use LinearEnsemble

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-linear-ensemble.md

Shows how to instantiate LinearEnsemble with in_features=2, out_features=3, k=4, and pass a random tensor of shape (5,4,2) to get output shape (5,4,3).

```python
import torch
import tabm

m = tabm.LinearEnsemble(2, 3, k=4)
x = torch.randn(5, 4, 2)
print(m(x).shape)  # torch.Size([5, 4, 3])
```

--------------------------------

### Example: using batchensemble_linear_layers_

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-functions.md

Example usage of batchensemble_linear_layers_ on a sequential backbone. It replaces Linear layers with LinearBatchEnsemble using k=32 and TabM-style initialization.

```python
import torch.nn as nn
import tabm

backbone = nn.Sequential(nn.Linear(24, 256), nn.ReLU(), nn.Linear(256, 256))
tabm.batchensemble_linear_layers_(
    backbone, k=32,
    scaling_init=('random-signs', 'ones'),   # TabM-style initialization
)
assert isinstance(backbone[0], tabm.LinearBatchEnsemble)
```

--------------------------------

### Import required libraries

Source: https://github.com/yandex-research/tabm/blob/main/example.ipynb

Import all necessary Python libraries for the example.

```python
import math
import random
from copy import deepcopy
from typing import Any, Literal, NamedTuple, Optional

import numpy as np
import rtdl_num_embeddings  # https://github.com/yandex-research/rtdl-num-embeddings
import scipy.special
import sklearn.datasets
import sklearn.metrics
import sklearn.model_selection
import sklearn.preprocessing
import tabm
import torch
import torch.nn as nn
import torch.optim
from torch import Tensor
```

--------------------------------

### Example: from_linear usage

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-linear-ensemble.md

Shows how to create a LinearEnsemble from an nn.Linear layer with k=16, and asserts that in_features, out_features, and k are preserved.

```python
import torch.nn as nn
import tabm

linear = nn.Linear(4, 8)
ens = tabm.LinearEnsemble.from_linear(linear, k=16)
assert ens.in_features == 4 and ens.out_features == 8 and ens.k == 16
```

--------------------------------

### Install TabM via pip

Source: https://github.com/yandex-research/tabm/blob/main/README.md

Use this command to install the TabM package in your environment.

```bash
pip install tabm
```

--------------------------------

### Basic usage example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-view.md

Shows how to use EnsembleView in a sequential model with k=4, expanding a (3,8) input to (3,4,8).

```python
import torch
import torch.nn as nn
import tabm

model = nn.Sequential(
    tabm.EnsembleView(k=4),
    nn.Linear(8, 8),          # operates on (B, 4, 8) with fully shared weights
)
x = torch.randn(3, 8)
y = model(x)
assert y.shape == (3, 4, 8)
```

--------------------------------

### TabM.make usage example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-model.md

Example showing how to create TabM models with default hyperparameters, one custom hyperparameter, and with feature embeddings. The final assertion checks the output shape.

```python
from tabm import TabM
from rtdl_num_embeddings import LinearReLUEmbeddings

# Default hyperparameters.
model = TabM.make(n_num_features=16, d_out=1)

# One custom hyperparameter, the rest defaulted.
model = TabM.make(n_num_features=16, d_out=1, n_blocks=2)

# With (shared) feature embeddings.
model = TabM.make(
    n_num_features=16,
    num_embeddings=LinearReLUEmbeddings(16),
    d_out=1,
)
assert model(torch.randn(8, 16)).shape == (8, model.k, 1)
```

--------------------------------

### LinearBatchEnsemble example with tuple scaling init and error handling

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-linear-batch-ensemble.md

Example usage of LinearBatchEnsemble with tuple scaling init and error handling. Creates a model with in_features=4, out_features=5, k=3, and different scaling distributions for r and s. Demonstrates the expected output shape and the ValueError raised for 2D input.

```python
import torch
import tabm

m = tabm.LinearBatchEnsemble(
    4, 5, k=3,
    scaling_init=('random-signs', 'ones'),   # different dists for r and s
)
y = m(torch.randn(2, 3, 4))
assert y.shape == (2, 3, 5)

try:
    m(torch.randn(2, 4))                     # 2D input
except ValueError as err:
    print(err)  # "The input must have 3 dimensions, however: x.ndim=2"

```

--------------------------------

### Initialize and use MLPBackboneBatchEnsemble

Source: https://github.com/yandex-research/tabm/blob/main/README.md

Example demonstrating the creation of a sequential model using EnsembleView and MLPBackboneBatchEnsemble to process input tensors.

```python
import tabm
import torch
import torch.nn as nn

d_in = 24
d_out = 1
k = 32
model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.MLPBackboneBatchEnsemble(
        d_in=d_in,
        n_blocks=3,
        d_block=512,
        dropout=0.1,
        k=k,
        tabm_init=True,
        scaling_init='normal',
        start_scaling_init_chunks=None,
    ),
    tabm.LinearEnsemble(512, d_out, k=k)
)
x = torch.randn(batch_size, d_in)
y_pred = model(x)
assert y_pred.shape == (batch_size, k, d_out)
```

--------------------------------

### Model and Optimizer Setup with Embedding Options

Source: https://github.com/yandex-research/tabm/blob/main/example.ipynb

Shows how to create a TabM model with different embedding options (none, LinearReLU, Periodic, PiecewiseLinear) and set up the optimizer with AdamW. Includes optional torch.compile with a note about mode='reduce-overhead' causing issues.

```python
# No embeddings.
num_embeddings = None

# Simple embeddings.
num_embeddings = rtdl_num_embeddings.LinearReLUEmbeddings(n_num_features)

# Periodic embeddings.
num_embeddings = rtdl_num_embeddings.PeriodicEmbeddings(n_num_features, lite=False)

# Piecewise-linear embeddings.
num_embeddings = rtdl_num_embeddings.PiecewiseLinearEmbeddings(
    rtdl_num_embeddings.compute_bins(data['train']['x_num'], n_bins=48),
    d_embedding=16,
    activation=False,
    version='B',
)

model = tabm.TabM.make(
    n_num_features=n_num_features,
    cat_cardinalities=cat_cardinalities,
    d_out=1 if n_classes is None else n_classes,
    num_embeddings=num_embeddings,
).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=3e-4)
gradient_clipping_norm: Optional[float] = 1.0

if compile_model:
    # NOTE
    # `torch.compile(model, mode="reduce-overhead")` caused issues during training,
    # so the `mode` argument is not used.
    model = torch.compile(model)
    evaluation_mode = torch.no_grad
else:
    evaluation_mode = torch.inference_mode
```

--------------------------------

### Example: Converting a backbone to an ensemble with k=32

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-functions.md

Shows how to convert a simple MLP backbone into an ensemble of 32 independent MLPs using `ensemble_linear_layers_`. The example asserts that the first layer becomes a `LinearEnsemble` and that the final model outputs the expected shape (batch, k, output).

```python
import torch
import torch.nn as nn
import tabm

backbone = nn.Sequential(nn.Linear(24, 256), nn.ReLU(), nn.Linear(256, 256))
tabm.ensemble_linear_layers_(backbone, k=32)
assert isinstance(backbone[0], tabm.LinearEnsemble)

model = nn.Sequential(
    tabm.EnsembleView(k=32),
    backbone,
    tabm.LinearEnsemble(256, 1, k=32),
)
assert model(torch.randn(8, 24)).shape == (8, 32, 1)
```

--------------------------------

### List public API items using uv

Source: https://github.com/yandex-research/tabm/blob/main/README.md

Use this command to list all available items in the tabm package without needing to clone or install the repository locally.

```bash
uv run --no-project --with tabm python -c """
import tabm

for x in sorted(
    x
    for x in dir(tabm)
    if getattr(getattr(tabm, x), '__module__', None) == 'tabm'
    and not x.startswith('_')
):
    print(x)
"""
```

--------------------------------

### Example: using ensemble_batchnorm1d_layers_

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-functions.md

Example usage of ensemble_batchnorm1d_layers_ on a sequential model. It replaces the BatchNorm1d layer with BatchNorm1dEnsemble using k=5.

```python
import torch.nn as nn
import tabm

model = nn.Sequential(nn.Linear(2, 3), nn.BatchNorm1d(3), nn.ReLU())
tabm.ensemble_batchnorm1d_layers_(model, k=5)
assert isinstance(model[1], tabm.BatchNorm1dEnsemble)
```

--------------------------------

### Tune hyperparameters with bin/tune.py

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Runs hyperparameter tuning for bin/model.py. The --continue flag resumes an interrupted run. This example reproduces the TabM tuning on California Housing and takes about one hour on an NVIDIA A100.

```bash
mkdir -p exp/reproduce/tabm/california
cp exp/tabm/california/0-tuning.toml exp/reproduce/tabm/california
python bin/tune.py exp/reproduce/tabm/california/0-tuning.toml --continue
```

--------------------------------

### LayerNormEnsemble usage example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-layernorm-ensemble.md

Shows how to instantiate LayerNormEnsemble with different normalized_shape values and verify output shapes. The first example uses a scalar normalized_shape (2) with k=3, and the second uses a tuple (2,3) with k=4.

```python
import torch
import tabm

m = tabm.LayerNormEnsemble(2, k=3)
x = torch.randn(4, 3, 2)
print(m(x).shape)  # torch.Size([4, 3, 2])

m = tabm.LayerNormEnsemble((2, 3), k=4)
x = torch.randn(5, 6, 4, 2, 3)
print(m(x).shape)  # torch.Size([5, 6, 4, 2, 3])
```

--------------------------------

### Device Setup, Data Conversion, AMP, and torch.compile

Source: https://github.com/yandex-research/tabm/blob/main/example.ipynb

Configures the device (CUDA if available), converts data to tensors, sets up automatic mixed precision (AMP) with bfloat16 or float16, and optionally enables torch.compile. The AMP is disabled by default (amp_enabled=False) and the output shows the device, AMP status, and compile status.

```python
# Device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Convert data to tensors
data = {
    part: {k: torch.as_tensor(v, device=device) for k, v in data_numpy[part].items()}
    for part in data_numpy
}
Y_train = torch.as_tensor(Y_train, device=device)
if task_type == 'regression':
    for part in data:
        data[part]['y'] = data[part]['y'].float()
    Y_train = Y_train.float()

# Automatic mixed precision (AMP)
# torch.float16 is implemented for completeness,
# but it was not tested in the project,
# so torch.bfloat16 is used by default.
amp_dtype = (
    torch.bfloat16
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    else torch.float16
    if torch.cuda.is_available()
    else None
)
# Changing False to True can speed up training
# of large enough models on compatible hardware.
amp_enabled = False and amp_dtype is not None
grad_scaler = torch.cuda.amp.GradScaler() if amp_dtype is torch.float16 else None  # type: ignore

# torch.compile
compile_model = False

# fmt: off
print(f'Device:        {device.type.upper()}')
print(f'AMP:           {amp_enabled}{f" ({amp_dtype})"if amp_enabled else ""}')
print(f'torch.compile: {compile_model}')
# fmt: on
```

```text
Output:
Device:        CUDA
AMP:           False
torch.compile: False
```

--------------------------------

### make_tabm_backbone

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-make-tabm-backbone.md

Documentation for the make_tabm_backbone function, including parameters, return values, validation rules, and a usage example.

```APIDOC
## make_tabm_backbone

### Description
Creates a TabM backbone with the specified architecture type and returns an `MLPBackboneEnsembleBase` instance. The concrete runtime type depends on the `arch_type` argument.

### Method
Function call

### Signature
```python
tabm.make_tabm_backbone(
    d_in: int,
    n_blocks: int,
    d_block: int,
    dropout: float,
    k: int,
    arch_type: str,
    start_scaling_init: Optional[str] = None,
    start_scaling_init_chunks: Optional[List[int]] = None,
    **kwargs
) -> MLPBackboneEnsembleBase
```

### Parameters
- **d_in** (int) - Required - Input dimensionality.
- **n_blocks** (int) - Required - Number of blocks.
- **d_block** (int) - Required - Block dimensionality.
- **dropout** (float) - Required - Dropout rate.
- **k** (int) - Required - Number of submodels.
- **arch_type** (str) - Required - One of `'tabm'`, `'tabm-mini'`, or `'tabm-packed'`.
- **start_scaling_init** (Optional[str]) - Optional - Initialization for scaling; must be `None` for `'tabm-packed'` and not `None` for other types.
- **start_scaling_init_chunks** (Optional[List[int]]) - Optional - Chunk sizes for scaling initialization; must be `None` if `start_scaling_init` is `None`.
- **kwargs** - Optional - Additional arguments passed to the backbone constructor.

### Return Value
`MLPBackboneEnsembleBase` — the created backbone, whose concrete runtime type depends on `arch_type`:

| `arch_type` | Returned instance |
|---|---|
| `'tabm'` | `MLPBackboneBatchEnsemble` (`tabm_init=True`, `scaling_init=start_scaling_init`) |
| `'tabm-mini'` | `MLPBackboneMiniEnsemble` (`affine_scaling_init=start_scaling_init`, `affine_bias=False`) |
| `'tabm-packed'` | `MLPBackboneEnsemble` (independent per-submodel layers) |

### Validation Rules and Throws

| Error type | Condition |
|---|---|
| `ValueError` | `arch_type == 'tabm-packed'` and `start_scaling_init is not None` ("When arch_type='tabm-packed', start_scaling_init must be None"). |
| `ValueError` | `arch_type != 'tabm-packed'` and `start_scaling_init is None` ("When arch_type='tabm', start_scaling_init must not be None"). |
| `ValueError` | `start_scaling_init is None` and `start_scaling_init_chunks is not None`. |
| `ValueError` | `arch_type` not in `{'tabm', 'tabm-mini', 'tabm-packed'}` ("Unknown arch_type='...'"). |
| `ValueError` | `d_in`/`n_blocks`/`d_block` not positive (propagated from `_MLPBackboneBase`). |
| `ValueError` | `activation` not an attribute of `torch.nn` (propagated). |
| `UserWarning` | `arch_type` is `'tabm'` or `'tabm-mini'` and `start_scaling_init_chunks is None` — "start_scaling_init_chunks is not provided, which may lead to suboptimal performance in some cases" (`warnings.warn`). |

### Example (basic)
```python
import torch
import tabm

backbone = tabm.make_tabm_backbone(
    d_in=24,
    n_blocks=3,
    d_block=512,
    dropout=0.1,
    k=32,
    arch_type='tabm',
    start_scaling_init='random-signs',   # or 'normal' when embeddings precede the backbone
    start_scaling_init_chunks=[1] * 24,  # one chunk per feature
)
out = backbone(torch.randn(8, 32, 24))   # (8, 32, 512)
```
```

--------------------------------

### TabM.forward backbone outputs example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-model.md

Shows how to obtain backbone features by setting d_out=None. The model is created with n_num_features=8 and d_out=None, and the output shape is (16, 32, 512).

```python
model = TabM.make(n_num_features=8, d_out=None)        # expose backbone features
out = model(torch.randn(16, 8))
print(out.shape)                                       # (16, 32, 512)
```

--------------------------------

### 3D pass-through and warning example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-view.md

Demonstrates 3D input pass-through, the warning in eval mode, and the ValueError for mismatched penultimate dimension.

```python
import warnings
import torch
import tabm

m = tabm.EnsembleView(k=2)
x3d = torch.randn(3, 2, 4)
assert m(x3d) is x3d          # propagated as-is, no copy

m.eval()
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter('always')
    m(x3d)
    print([str(x.message) for x in w])
    # ["When training=False, the input should usually have the shape (batch_size, d), ..."]

try:
    m(torch.randn(3, 5, 4))   # penultimate dim != k
except ValueError as err:
    print(err)  # "The penultimate input dimension must be equal to k=2, however: x.shape[-2]=5"
```

--------------------------------

### Create a basic TabM model

Source: https://github.com/yandex-research/tabm/blob/main/README.md

Initialize a basic TabM model without feature embeddings. This example demonstrates setting up numerical and categorical features and performing a forward pass.

```python
import torch
from tabm import TabM

# >>> Common setup for all subsequent sections.
d_out = 1  # For example, one regression task.
batch_size = 256

# The dataset has 24 numerical (continuous) features.
n_num_features = 24

# The dataset has 2 categorical features.
# The first categorical feature has 3 unique categories.
# The second categorical feature has 7 unique categories.
cat_cardinalities = [3, 7]
# <<<

model = TabM.make(
    n_num_features=n_num_features,
    cat_cardinalities=cat_cardinalities,  # One-hot encoding will be used.
    d_out=d_out,
)
x_num = torch.randn(batch_size, n_num_features)
x_cat = torch.column_stack([
    # The i-th categorical features must take values in range(0, cat_cardinalities[i]).
    torch.randint(0, c, (batch_size,)) for c in cat_cardinalities
])
y_pred = model(x_num, x_cat)

# TabM represents an ensemble of k models, hence k predictions per object.
assert y_pred.shape == (batch_size, model.k, d_out)
```

--------------------------------

### TabM.forward basic and 3D batched input example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-model.md

Shows how to call the model with 2D and 3D numerical input. The model is created with n_num_features=8 and d_out=1, resulting in k=32 submodels. The first call uses a single batch shared by all submodels, producing shape (64, 32, 1). The second call provides k full-fledged batches with shape (64, 32, 8), also producing (64, 32, 1).

```python
import torch
from tabm import TabM, LinearEnsemble

model = TabM.make(n_num_features=8, d_out=1)          # k = 32

# (A) One batch shared by all submodels.
x_num = torch.randn(64, 8)
y_pred = model(x_num)                                  # (64, 32, 1)

# (B) k full-fledged batches: x_num has shape (64, 32, 8).
x_num = torch.randn(64, model.k, 8)
y_pred = model(x_num)                                  # (64, 32, 1)
```

--------------------------------

### TabM error handling examples

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-model.md

Shows how to catch ValueError and TypeError from invalid constructor arguments, and how to infer d_out by passing None. The output shapes and error messages are documented.

```python
from tabm import TabM

try:
    TabM(n_num_features=0, cat_cardinalities=[], d_out=1, k=32,
         start_scaling_init='random-signs')
except ValueError as err:
    print(err)  # "n_num_features=0 and cat_cardinalities=[], at the same time, ..."

# Passing a forbidden backbone argument is rejected:
try:
    TabM(n_num_features=16, d_out=1, k=32, start_scaling_init='normal', d_in=16)
except TypeError as err:
    print(err)  # "The argument d_in must not be passed explicitly"

# Infer d_out from the backbone by passing None:
model = TabM(n_num_features=16, d_out=None, k=16, start_scaling_init='normal')
print(model(torch.randn(4, 16)).shape)  # (4, 16, 512) — d_block of the default backbone

```

--------------------------------

### Quick test with Pixi and without

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Run a quick test to verify the environment. The first run may be slow, and results are not representative.

```shell
# Pixi with GPU
pixi run -e cuda python bin/model.py exp/debug/0.toml --force

# Pixi without GPU
pixi run python bin/model.py exp/debug/0.toml --force

# Without Pixi
python bin/model.py exp/debug/0.toml --force
```

--------------------------------

### Prepare TabReD benchmark

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Download the TabReD benchmark to local/tabred (requires Kaggle account), then run this command to prepare the data.

```shell
python tools/prepare_tabred.py local/tabred data
```

--------------------------------

### MLPBackboneEnsemble usage example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-backbones.md

Example of creating an MLPBackboneEnsemble with d_in=24, n_blocks=2, d_block=512, dropout=0.1, and k=32, then passing a random tensor of shape (8, 32, 24) to get output of shape (8, 32, 512).

```python
import torch
import tabm

model = tabm.MLPBackboneEnsemble(d_in=24, n_blocks=2, d_block=512, dropout=0.1, k=32)
out = model(torch.randn(8, 32, 24))         # (8, 32, 512)
```

--------------------------------

### init_scaling_ chunked initialization

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-initialization.md

Demonstrates chunked initialization and handling of invalid configurations.

```python
import torch
import tabm

x = torch.empty(4, 5)
tabm.init_scaling_(x, 'normal', chunks=[2, 3])
# Each row is split into two chunks of sizes 2 and 3;
# all values within one chunk share one sampled value.

try:
    tabm.init_scaling_(x, 'ones', chunks=[2, 3])
except ValueError as err:
    print(err)          # "When distribution='ones', chunks must be None"

try:
    tabm.init_scaling_(x, 'normal', chunks=[2, 2])
except ValueError as err:
    print(err)          # "The tensor shape and chunks are incompatible: 5 != 4"
```

--------------------------------

### Train a model with bin/model.py

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Creates a directory, copies a TOML config, and runs a single training run. The output directory is created next to the config, and report.json contains the main results.

```bash
mkdir -p exp/reproduce/train-once
cp exp/tabm/california/0-evaluation/0.toml exp/reproduce/train-once/
python bin/model.py exp/reproduce/train-once/0.toml
```

--------------------------------

### ElementwiseAffine usage example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-elementwise-affine.md

Example demonstrating the initialization and application of ElementwiseAffine to a tensor.

```python
import torch
import tabm

m = tabm.ElementwiseAffine(
    (3, 4),
    scaling_init='normal',
    scaling_init_chunks=None,
    bias=True,
)
x = torch.randn((1, 2, 3, 4))
print(m(x).shape)       # torch.Size([1, 2, 3, 4])
print(m.weight.shape)   # torch.Size([3, 4])
```

--------------------------------

### MLPBackbone usage example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-backbones.md

Example of creating an MLPBackbone and applying it to a random tensor.

```python
import torch
import tabm

backbone = tabm.MLPBackbone(d_in=24, n_blocks=3, d_block=512, dropout=0.1)
out = backbone(torch.randn(8, 24))          # (8, 512)
```

--------------------------------

### ElementwiseAffine error handling example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-elementwise-affine.md

Example showing successful application and a ValueError when input shapes do not match.

```python
import torch
import tabm

m = tabm.ElementwiseAffine((2, 3), bias=False, scaling_init='ones')
print(m(torch.randn(4, 2, 3)).shape)            # (4, 2, 3)

try:
    m(torch.randn(4, 3, 2))                      # wrong trailing shape
except ValueError as err:
    print(err)  # "The input must have the shape (*, 2, 3), however: x.shape=torch.Size([4, 3, 2])"
```

--------------------------------

### Initialize and run TabM model

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm.md

Demonstrates the end-to-end usage of TabM, including model creation, input generation, and inference.

```python
import torch
from tabm import TabM

model = TabM.make(n_num_features=24, cat_cardinalities=[3, 7], d_out=1)
x_num = torch.randn(256, 24)
x_cat = torch.column_stack([
    torch.randint(0, c, (256,)) for c in [3, 7]
])
y_pred = model(x_num, x_cat)          # (256, k=32, 1)
y_pred.mean(dim=1)                    # ensemble prediction for inference
```

--------------------------------

### TabM.make

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/configuration.md

Documentation for the TabM.make constructor method, including all options with defaults and required arguments.

```APIDOC
## TabM.make

### Description
Creates a TabM model with the given data-related arguments and default model hyperparameters. Explicit keyword arguments override the computed defaults.

### Method
make

### Endpoint
TabM.make(**kwargs)

### Parameters
#### Keyword Arguments
- **n_num_features** (int) - Required - The number of numerical (continuous) features.
- **cat_cardinalities** (Optional[list[int]]) - Optional - Cardinalities of the categorical features; the i-th feature must take values in range(0, cat_cardinalities[i]).
- **d_out** (Optional[int]) - Required - Output size; None returns backbone features (B, k, d_block).
- **num_embeddings** (Optional[_NumEmbeddings]) - Optional - Feature embeddings from rtdl_num_embeddings (LinearReLUEmbeddings, PiecewiseLinearEmbeddings with version="B", PeriodicEmbeddings); fully shared between the k MLPs.
- **n_blocks** (int) - Optional - The number of blocks (depth) of the backbone. Default: 2 if num_embeddings else 3.
- **d_block** (int) - Optional - Latent representation size (width). Default: 512.
- **dropout** (float) - Optional - Dropout rate. Default: 0.1.
- **activation** (str) - Optional - Activation; must be a valid attribute of torch.nn. Default: 'ReLU'.
- **k** (int) - Optional - The number of submodels (ensemble size). Default: 32.
- **arch_type** (TabMArchitectureType) - Optional - Architecture variant ('tabm', 'tabm-mini', 'tabm-packed'). Default: 'tabm'.
- **start_scaling_init** (Optional[ScalingRandomInitialization]) - Optional - Initialization of the very first scaling. Default: None for 'tabm-packed'; 'normal' if num_embeddings provided; else 'random-signs'.
- **start_scaling_init_chunks** (Optional[list[int]]) - Optional - Initialization chunks of the very first scaling; must sum exactly to d_in. Passing it to TabM.__init__/TabM.make is forbidden — TabM infers it as d_features (one chunk per feature representation).

### Response
Returns a TabM model instance.

### Notes
- Hyperparameter semantics from the README (guidance, not code):
  - arch_type='tabm' is the default and expected best; 'tabm-mini' may train/infer faster but needs more precision tuning d_block/n_blocks per k; 'tabm-packed' is implemented for completeness only.
  - For a given depth/width, increasing k helps up to a threshold. Consider increasing d_block or n_blocks when increasing k. n_blocks=1 should usually be avoided unless the tuning budget is high. Avoid tuning k together with other hyperparameters.
  - 'normal' scaling init is recommended when "non-trivial" modules (e.g. embeddings) precede the backbone; 'random-signs' otherwise (informal heuristic, README "Initialization").
```

--------------------------------

### BatchNorm1dEnsemble.from_batchnorm1d example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-batchnorm1d-ensemble.md

Converts a standard BatchNorm1d with 3 features into an ensemble with k=7 and verifies the ensemble size.

```python
import torch.nn as nn
import tabm

bn = nn.BatchNorm1d(3)
ens = tabm.BatchNorm1dEnsemble.from_batchnorm1d(bn, k=7)
assert ens.k == 7
```

--------------------------------

### MLPBackboneEnsembleBase.get_original_output_shape

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-backbones.md

Abstract method to get the output shape of one underlying MLP, without batch dimensions.

```APIDOC
## MLPBackboneEnsembleBase.get_original_output_shape

### Description
Abstract method to get the output shape of one underlying MLP, without batch dimensions.

### Method
Abstract method

### Signature
```python
@abc.abstractmethod
def get_original_output_shape(self) -> torch.Size:
```

### Returns
- **torch.Size** - The output shape of one underlying MLP, without batch dimensions.

```

--------------------------------

### Create environment with Micromamba

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Create and activate the environment using Micromamba.

```shell
micromamba create --file environment.yaml
micromamba activate tabm
```

--------------------------------

### Create environment with Mamba

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Create and activate the environment using Mamba.

```shell
mamba create --file environment.yaml
mamba activate tabm
```

--------------------------------

### Automate tuning, evaluation, and ensembling with bin/go.py

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Runs hyperparameter tuning, evaluation, and ensembling in a single command. The --continue flag resumes an interrupted run.

```bash
mkdir -p exp/reproduce/tabm-go/california
cp exp/tabm/california/0-tuning.toml exp/reproduce/tabm-go/california

python bin/go.py exp/reproduce/tabm-go/california/0-tuning --continue
```

--------------------------------

### BatchNorm1dEnsemble forward example

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-batchnorm1d-ensemble.md

Demonstrates forward in train and eval modes, and catches the ValueError when the penultimate dimension does not match k.

```python
import torch
import tabm

m = tabm.BatchNorm1dEnsemble(3, k=4, affine=True)
m.train()
y = m(torch.randn(16, 4, 3))
assert y.shape == (16, 4, 3)

# In eval mode, running statistics are used.
m.eval()
y = m(torch.randn(4, 4, 3))

try:
    m(torch.randn(4, 5, 3))       # penultimate dim != k
except ValueError as err:
    print(err)
```

--------------------------------

### Create environment with Conda

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

Create and activate the environment using Conda.

```shell
conda create --file environment.yaml -n tabm
conda activate tabm
```

--------------------------------

### Order of layers in TabM: create k representations before linear mixing

Source: https://github.com/yandex-research/tabm/blob/main/README.md

Shows correct and incorrect layer ordering for TabM models. Correct examples use EnsembleView before linear mixing, with ElementwiseAffine or LinearBatchEnsemble to create diverse representations. Incorrect examples mix features before ensembling or use a shared linear layer. This is a key implementation guideline directly derived from the paper's Section 3.3.

```python
d_in = 24
d = 512
k = 16

# GOOD: collectively, the first two modules create k
# different object representations before the first
# linear layer.
scaling_init = "normal"  # or "random-signs"
good_model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.ElementwiseAffine((k, d_in), bias=False, scaling_init=scaling_init),
    nn.Linear(d_in, d),
    ...
)

# GOOD: internally, LinearBatchEnsemble starts with a
# non-shared elementwise scaling, which diversifies
# the k object representations before the
# linear transformation.
scaling_init = "normal"  # or "random-signs"
                         # or ("normal", "ones")
                         # or ("random-signs", "ones")
good_model = nn.Sequential(
    tabm.EnsembleView(k=k),
    tabm.LinearBatchEnsemble(d_in, d, k=k, scaling_init=scaling_init),
    ...
)

# BAD: the tabular features are mixed before
# the ensemble starts.
bad_model = nn.Sequential(
    nn.Linear(d_in, d),
    tabm.EnsembleView(k=k),
    nn.ReLU(),
    ...
)

# BAD: the k representations are created before the
# first linear transformations, but these representations
# are not different. Mathematically, the below snippet is
# equivalent to the previous one.
bad_model = nn.Sequential(
    tabm.EnsembleView(k=k),
    nn.Linear(d_in, d),
    nn.ReLU(),
    ...
)
```

--------------------------------

### init_scaling_ basic usage

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-initialization.md

Basic usage of init_scaling_ with different distribution types.

```python
import torch
import tabm

w = torch.empty(4, 5)
tabm.init_scaling_(w, 'random-signs')
print(w.unique().tolist())  # [-1.0, 1.0]

tabm.init_scaling_(w, 'normal')
tabm.init_scaling_(w, 'ones')   # all ones
```

--------------------------------

### MLPBackboneEnsembleBase abstract method get_original_output_shape

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-backbones.md

Abstract method to get the output shape of one underlying MLP, without batch dimensions.

```python
@abc.abstractmethod
def get_original_output_shape(self) -> torch.Size:
```

--------------------------------

### Example: forward with error handling

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-linear-ensemble.md

Shows forward pass with correct input shape and a try-except for wrong last dimension, printing the ValueError message.

```python
import torch
import tabm

m = tabm.LinearEnsemble(4, 8, bias=True, k=16)
y = m(torch.randn(32, 16, 4))
assert y.shape == (32, 16, 8)

try:
    m(torch.randn(32, 16, 5))          # wrong last dimension
except ValueError as err:
    print(err)  # "The last input dimension must be equal to 4, however: 5"
```

--------------------------------

### Evaluate a manually composed config with bin/evaluate.py

Source: https://github.com/yandex-research/tabm/blob/main/paper/README.md

To evaluate a custom config, create a directory ending with -evaluation and place the config as 0.toml inside. Then pass the directory to bin/evaluate.py with the --function flag.

```bash
# The config is stored at exp/<any/path>/0-evaluation/0.toml
python bin/evaluate.py exp/<any/path>/0-evaluation --function "bin.model.main"
```

--------------------------------

### LayerNormEnsemble forward example with error handling

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-layernorm-ensemble.md

Demonstrates forward pass with correct shape and catches ValueError for incorrect trailing shape. The error message is shown.

```python
import torch
import tabm

m = tabm.LayerNormEnsemble((2, 3), k=4, elementwise_affine=True)
y = m(torch.randn(5, 4, 2, 3))
assert y.shape == (5, 4, 2, 3)

try:
    m(torch.randn(5, 4, 3, 2))       # wrong trailing shape
except ValueError as err:
    print(err)  # "The input must have the shape (*, 4, 2, 3), however: x.shape=..."
```

--------------------------------

### Using include and exclude parameters with ensemble_linear_layers_

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-functions.md

Demonstrates the `include` and `exclude` parameters for `ensemble_linear_layers_`. The first part uses `include` to replace only the first linear layer, the second uses `exclude` to replace all but the first, and the third shows that providing both raises a `ValueError` with the message "include and exclude cannot be provided simultaneously".

```python
import torch.nn as nn
import tabm

model = nn.Sequential(nn.Linear(2, 3), nn.ReLU(), nn.Linear(3, 4))
tabm.ensemble_linear_layers_(model, include=[model[0]], k=5)
assert isinstance(model[0], tabm.LinearEnsemble)
assert not isinstance(model[2], tabm.LinearEnsemble)

model = nn.Sequential(nn.Linear(2, 3), nn.ReLU(), nn.Linear(3, 4))
tabm.ensemble_linear_layers_(model, exclude=[model[0]], k=5)
assert not isinstance(model[0], tabm.LinearEnsemble)
assert isinstance(model[2], tabm.LinearEnsemble)

try:
    tabm.ensemble_linear_layers_(model, include=[model[0]], exclude=[model[2]], k=5)
except ValueError as err:
    print(err)  # "include and exclude cannot be provided simultaneously"
```

--------------------------------

### Create TabM backbone with make_tabm_backbone

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-make-tabm-backbone.md

Creates a TabM backbone with the 'tabm' architecture and applies it to a random input. Requires d_in, n_blocks, d_block, dropout, k, arch_type, start_scaling_init, and start_scaling_init_chunks. The output shape is (8, 32, 512).

```python
import torch
import tabm

backbone = tabm.make_tabm_backbone(
    d_in=24,
    n_blocks=3,
    d_block=512,
    dropout=0.1,
    k=32,
    arch_type='tabm',
    start_scaling_init='random-signs',   # or 'normal' when embeddings precede the backbone
    start_scaling_init_chunks=[1] * 24,  # one chunk per feature
)
out = backbone(torch.randn(8, 32, 24))   # (8, 32, 512)
```

--------------------------------

### EnsembleView.__init__

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-view.md

Constructor for the EnsembleView module. Initializes the ensemble size k.

```APIDOC
## EnsembleView.__init__

### Description
Initializes the EnsembleView module with the specified ensemble size.

### Method
Constructor

### Signature
```python
def __init__(self, *, k: int) -> None:
```

### Parameters
- **k** (int) - Required (keyword-only) - The ensemble size.

```

--------------------------------

### Replace LayerNorm submodules with LayerNormEnsemble

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-ensemble-functions.md

Replaces all `torch.nn.LayerNorm` submodules in a model with `LayerNormEnsemble` using `LayerNormEnsemble.from_layernorm`. The `kwargs` (e.g., `k`) are passed to `LayerNormEnsemble`. The example creates a simple sequential model and asserts the replacement.

```python
import torch.nn as nn
import tabm

model = nn.Sequential(nn.Linear(2, 3), nn.LayerNorm(3))
tabm.ensemble_layernorm_layers_(model, k=5)
assert isinstance(model[1], tabm.LayerNormEnsemble)
```

--------------------------------

### NLinear.__init__

Source: https://github.com/yandex-research/tabm/blob/main/_autodocs/api-reference/tabm-reference-module.md

Initializes NLinear, a module that applies n linear layers in parallel to n disjoint input matrices. It is a simplified version of delu.nn.NLinear supporting 3D inputs with one batch dimension.

```APIDOC
## NLinear.__init__

### Description
Initializes NLinear, a module that applies n linear layers in parallel to n disjoint input matrices. It is a simplified version of delu.nn.NLinear supporting 3D inputs with one batch dimension.

### Method
Constructor

### Signature
```python
def __init__(
    self, n: int, in_features: int, out_features: int, bias: bool = True
) -> None:
```

### Parameters
- **n** (int) - Required - Number of linear layers.
- **in_features** (int) - Required - Number of input features for each layer.
- **out_features** (int) - Required - Number of output features for each layer.
- **bias** (bool) - Optional - Whether to include bias terms. Default: True.

### Attributes
- **weight** (Tensor) - Shape (n, in_features, out_features).
- **bias** (Tensor) - Optional, shape (n, out_features).

### Input/Output
- Input: (B, N, in_features)
- Output: (B, N, out_features)

### Usage
Used by arch_type='tabm-packed'.
```
Tokens:
Raw