KumoRelational

class sdm.models.KumoRelational(task: Task | str | Iterable[Task | str] | None = None, pretrained: bool = True, device: device | str | None = None)

Bases: ICLModel

An adapted version of the relational foundation model from the “KumoRFM-2: Scaling Foundation Models for Relational Learning” paper.

KumoRelational extends the in-context learning structure of tabular foundation models from single tables to relational, multi-table inputs. It processes task rows together with one or more related tables, avoiding manual flattening of relational data into a single table.

This implementation follows the high-level KumoRFM-2 architecture. It consists of three stages:

• Intra-table row embeddings: Each related table is embedded independently with a TabICLv2-style row embedding stack. Target information is injected by distributing context targets over the relational graph.

• Inter-table message passing: A schema-agnostic GNN exchanges the row embeddings along pre-defined relationships for num_hops rounds, before it reads out the rows linked to the task table.

• In-context learning over samples: The readout embeddings for context and query task rows are processed by a TabICLv2-style dataset-level ICL block. Context rows carry target information, while query rows attend to the labeled context to produce class logits or regression quantiles.

from sdm import RelatedTables, TableTensor
from sdm.models import KumoRelational

task_table = TableTensor.from_columns(
    {"user_id": [0, 1, 2, 3], "churn": [True, False, True, False]},
    stypes={"user_id": "id", "churn": "categorical"},
    device="cuda",
)

related_tables = RelatedTables(
    tables={
        "users": TableTensor.from_columns(
            {"user_id": [0, 1, 2, 3], "age": [42, 23, 31, 26]},
            stypes={"user_id": "id", "age": "numerical"},
            device="cuda",
        ),
        "orders": TableTensor.from_columns(
            {
                "user_id": [0, 0, 1, 3, 3, 3],
                "amount": [9.99, 4.99, 12.99, 7.99, 3.99, 5.99],
            },
            stypes={"user_id": "id", "amount": "numerical"},
            device="cuda",
        ),
    },
    relationships=[{
        "left_table": "orders",
        "left_columns": "user_id",
        "right_table": "users",
        "right_columns": "user_id",
    }],
    task_links=[{
        "task_columns": "user_id",
        "table": "users",
        "table_columns": "user_id",
    }],
)

x_context = task_table[:2].drop_columns("churn")
y_context = task_table[:2, "churn"]
x_query = task_table[2:].drop_columns("churn")

related_context_tables = related_tables.replace_tables({
    "users": related_tables.tables["users"][:2],
    "orders": related_tables.tables["orders"][:3],
})
related_query_tables = related_tables.replace_tables({
    "users": related_tables.tables["users"][2:],
    "orders": related_tables.tables["orders"][3:],
})

model = KumoRelational(device="cuda")

# Default in-context learning forward pass:
out = model(
    x_context=x_context,
    y_context=y_context,
    x_query=x_query,
    related_context_tables=related_context_tables,
    related_query_tables=related_query_tables,
    num_hops=1,
)

# Fit+Predict forward pass via key/value caching:
model.fit(x_context, y_context, related_context_tables, num_hops=1)
out = model.predict(x_query, related_query_tables)

Parameters:

• task (TaskLike | Iterable[TaskLike] | None) – The tasks to initialize. If None, all tasks supported by this model are initialized.

• pretrained (bool) – Whether to load the pretrained checkpoint.

• device (torch.device | str | None) – The device.

Capabilities

Supported Input Semantic Types	datetime, numerical
Supported Target Semantic Types	categorical, numerical
Supported Prediction Tasks	classification, regression
Multi-Target Support	❌
Related Table Support	✅

Default Recipe

Recipe(
  features=Sequential(
    TableDispatch(
      related=StypeDispatch(
        datetime=AddCalendarFields(
          fields=[
            'minute',
            'hour',
            'weekday',
            'day_of_month',
            'month',
          ],
          encoding='raw',
        ),
      ),
    ),
    StypeDispatch(
      categorical=Sequential(
        AlignCategories(sort_by='value'),
        ToNumerical(),
      ),
    ),
    StypeDispatch(
      numerical=Sequential(
        ImputeMean(),
        DropConstantColumns(),
        Standardize(eps=1e-06),
        Clip(-100.0, 100.0),
        Choice(
          Identity(),
          PowerTransform(),
          method='round_robin',
        ),
        ClipSigma(threshold=4.0),
        ShuffleColumns(method='latin'),
      ),
    ),
  ),
  target=Sequential(
    StypeDispatch(
      numerical=Standardize(eps=0.0),
      categorical=Sequential(
        AlignCategories(),
        ShuffleCategories(method='shift'),
      ),
    ),
  ),
  output=Sequential(
    TaskDispatch(
      regression=SortQuantiles(),
    ),
    AverageEstimators(),
    TaskDispatch(
      classification=Softmax(temperature=0.9),
    ),
  ),
)