Support padded contexts for Kumo Tabular

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:

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)

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)

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,