From 54f5a0f2697de807c82d54167ca382505ca3156d Mon Sep 17 00:00:00 2001 From: Valter Hudovernik Date: Thu, 29 Oct 2026 18:33:51 +0100 Subject: [PATCH] Support padded contexts in KumoTabular --- sdm/models/kumo/tabular/cell_embedding.py | 10 ++++-- sdm/models/kumo/tabular/icl.py | 38 ++++++++++++++++++----- sdm/models/kumo/tabular/model.py | 6 ++-- sdm/models/kumo/tabular/row_embedding.py | 14 +++++++-- 4 files changed, 53 insertions(+), 15 deletions(-) diff --git a/sdm/models/kumo/tabular/cell_embedding.py b/sdm/models/kumo/tabular/cell_embedding.py index 7529820b50..aa3e2d6bc9 100644 --- a/sdm/models/kumo/tabular/cell_embedding.py +++ b/sdm/models/kumo/tabular/cell_embedding.py @@ -42,7 +42,7 @@ def forward( # type: ignore x: Tensor, # [..., R, C] categorical_mask: Tensor, # [..., C] *, - train_size: int | None = None, + train_size: int | Tensor | None = None, num_columns: Tensor | None = None, # [...] cache: Cache | None = None, batch_size_limit: int | Literal["auto"] | None = None, @@ -54,7 +54,13 @@ def forward( # type: ignore impute_mean = cast(Tensor, cache["cell_impute_mean"]) else: assert train_size is not None - impute_mean = x[..., :train_size, :].nanmean(dim=-2, keepdim=True) + if isinstance(train_size, Tensor): + context = torch.arange(x.size(-2), device=x.device) + context = context < train_size[..., None] + observed = x.masked_fill(~context[..., None], torch.nan) + else: + observed = x[..., :train_size, :] + impute_mean = observed.nanmean(dim=-2, keepdim=True) impute_mean = impute_mean.masked_fill(impute_mean.isnan(), 0.0) if cache is not None and cache.is_recording: diff --git a/sdm/models/kumo/tabular/icl.py b/sdm/models/kumo/tabular/icl.py index efb75ff3f0..7f76b6dd86 100644 --- a/sdm/models/kumo/tabular/icl.py +++ b/sdm/models/kumo/tabular/icl.py @@ -63,8 +63,13 @@ def forward( y: Tensor, # [..., R_train] *, cache: Cache | None = None, - ) -> Tensor: # [..., R_test, out_channels] + train_size: Tensor | None = None, + ) -> Tensor: # [..., R or R_test, out_channels] R_train = y.size(-1) + context_mask = None + if train_size is not None: + context_mask = torch.arange(x.size(-2), device=x.device) + context_mask = context_mask < train_size[..., None] if y.numel() > 0: if self.y_emb is not None: @@ -73,6 +78,10 @@ def forward( assert self.y_lin is not None y_emb = self.y_lin(y.unsqueeze(-1)) # [..., R_train, D] + if context_mask is not None: + y_emb = y_emb.masked_fill( + ~context_mask[..., :R_train, None], 0 + ) x[..., :R_train, :] += y_emb.to(x.dtype) del y_emb @@ -82,15 +91,18 @@ def forward( if self.kv_heads is None or cache is not None: result = layer( - query=x[..., R_train:, :] if last_layer else x, + query=x[..., R_train:, :] + if last_layer and train_size is None + else x, key_value=( cast(KVCacheEntry, cache[cache_key]) if cache is not None and cache.is_replaying else x[..., :R_train, :] ), + seqused_key_value=train_size, return_key_value=cache is not None and cache.is_recording, out=None - if torch.is_grad_enabled() + if torch.is_grad_enabled() or train_size is not None else x[..., R_train:, :] if last_layer else x, @@ -112,11 +124,16 @@ def forward( continue x_context, (key, value) = layer( - query=x[..., :0, :] if last_layer else x[..., :R_train, :], + query=x + if train_size is not None + else x[..., :0, :] + if last_layer + else x[..., :R_train, :], key_value=x[..., :R_train, :], + seqused_key_value=train_size, return_key_value=True, out=None - if torch.is_grad_enabled() + if torch.is_grad_enabled() or train_size is not None else x[..., :0, :] if last_layer else x[..., :R_train, :], @@ -127,11 +144,16 @@ def forward( ) del key, value x_query = layer( - query=x[..., R_train:, :], + query=x if train_size is not None else x[..., R_train:, :], key_value=key_value, - out=None if torch.is_grad_enabled() else x[..., R_train:, :], + seqused_key_value=train_size, + out=None + if torch.is_grad_enabled() or train_size is not None + else x[..., R_train:, :], ) - if last_layer: + if context_mask is not None: + x = torch.where(context_mask[..., None], x_context, x_query) + elif last_layer: x = x_query elif torch.is_grad_enabled(): x = torch.cat([x_context, x_query], -2) diff --git a/sdm/models/kumo/tabular/model.py b/sdm/models/kumo/tabular/model.py index b946f963e3..b28cf4f55e 100644 --- a/sdm/models/kumo/tabular/model.py +++ b/sdm/models/kumo/tabular/model.py @@ -316,12 +316,14 @@ def forward_train( categorical_mask: Tensor, # [..., C] *, num_columns: Tensor | None = None, # [...] - ) -> Tensor: # [..., R_test, num_classes or num_quantiles] + train_size: Tensor | None = None, # [...], int32 lengths in padded y + ) -> Tensor: # [..., R or R_test, num_classes or num_quantiles] x = self.row_embedding( x=x, y=y, categorical_mask=categorical_mask, num_columns=num_columns, + train_size=train_size, ) x = self.row_project(x) - return self.icl_block(x=x, y=y) + return self.icl_block(x=x, y=y, train_size=train_size) diff --git a/sdm/models/kumo/tabular/row_embedding.py b/sdm/models/kumo/tabular/row_embedding.py index 6714c5ce5e..862c2af601 100644 --- a/sdm/models/kumo/tabular/row_embedding.py +++ b/sdm/models/kumo/tabular/row_embedding.py @@ -114,6 +114,7 @@ def forward( *, cache: Cache | None = None, num_columns: Tensor | None = None, + train_size: Tensor | None = None, ) -> Tensor: # [..., R, K * D] *B, R, C = x.size() @@ -127,11 +128,11 @@ def forward( ) buffer: Tensor | None = None - if torch.is_grad_enabled(): + if torch.is_grad_enabled() or train_size is not None: x = self.cell_embedding( x=x, categorical_mask=categorical_mask, - train_size=R_train, + train_size=R_train if train_size is None else train_size, num_columns=num_columns, cache=cache, ) # [..., R, C, D] @@ -147,7 +148,7 @@ def forward( x = self.cell_embedding( x=x, categorical_mask=categorical_mask, - train_size=R_train, + train_size=R_train if train_size is None else train_size, num_columns=num_columns, cache=cache, batch_size_limit="auto", @@ -160,6 +161,10 @@ def forward( else: assert self.y_lin is not None y_emb = self.y_lin(y.unsqueeze(-1)).unsqueeze(-2) + if train_size is not None: + context = torch.arange(R_train, device=y.device) + context = context < train_size[..., None] + y_emb = y_emb.masked_fill(~context[..., None, None], 0) x[..., :R_train, :, :] += y_emb.to(x.dtype) if buffer is not None: @@ -184,6 +189,9 @@ def forward( key_value=cast(KVCacheEntry, cache[key]) if cache is not None and cache.is_replaying else x[..., :R_train, :], + seqused_key_value=None + if train_size is None + else train_size.unsqueeze(-1), return_key_value=cache is not None and cache.is_recording, batch_size_limit="auto", out=None if buffer is None else x,