Add KumoAnomaly plumbing

sdm/models/base/icl.py

@@ -20,7 +20,6 @@
 )
 from sdm._inference import inference_mode
 from sdm._memory import set_pinned_max_round_threshold
-from sdm._warnings import warn_once
 from sdm.cache import Cache
 from sdm.models.base import Callback, Model
 from sdm.processing.execution import (
@@ -701,7 +700,8 @@ def _validate_context(
                 f"dimensions (got {tuple(x.size()[:-1])} and "
                 f"{tuple(y.size()[:-1])})"
             )
-        feature_stypes = x.active_stypes
+        self._validate_feature_stype(x)
+
         if x_context_only is not None:
             if x_context_only.size()[:-1] != y.size()[:-1]:
                 raise ValueError(
@@ -710,18 +710,8 @@ def _validate_context(
                     f"(got {tuple(x_context_only.size()[:-1])} and "
                     f"{tuple(y.size()[:-1])})"
                 )
-            feature_stypes = feature_stypes | x_context_only.active_stypes
-        invalid = feature_stypes - self.supported_feature_stypes - {Stype.id}
-        if len(invalid) > 0:
-            stypes = ", ".join(f"{str(stype)!r}" for stype in invalid)
-            warn_once(
-                key="model-unsupported-feature-stypes",
-                message=(
-                    f"{self.__class__.__name__!r} received unsupported "
-                    f"feature stypes {stypes}. Columns with unsupported "
-                    f"feature stypes will not be consumed by the model."
-                ),
-            )
+            self._validate_feature_stype(x_context_only)
+
         invalid = y.active_stypes - self.supported_target_stypes
         if len(invalid) > 0:
             stypes = ", ".join(f"{str(stype)!r}" for stype in invalid)
@@ -745,21 +735,8 @@ def _validate_context(
                     f"{self.__class__.__name__!r} does not support related "
                     f"tables"
                 )
-            for table_name, table in related_tables.tables.items():
-                invalid = table.active_stypes - self.supported_feature_stypes
-                invalid = invalid - {Stype.id}
-                if len(invalid) > 0:
-                    stypes = ", ".join(f"{str(stype)!r}" for stype in invalid)
-                    warn_once(
-                        key="model-unsupported-feature-stypes",
-                        message=(
-                            f"{self.__class__.__name__!r} received "
-                            f"unsupported feature stypes {stypes} in related "
-                            f"table {table_name!r}. Columns with unsupported "
-                            f"feature stypes will not be consumed by the "
-                            f"model."
-                        ),
-                    )
+            for table in related_tables.tables.values():
+                self._validate_feature_stype(table)
 
     def _validate_query(
         self,

sdm/models/base/model.py

@@ -8,7 +8,8 @@
 import torch
 from torch import Tensor
 
-from sdm import Recipe, Stype
+from sdm import Recipe, Stype, TableTensor
+from sdm._warnings import warn_once
 from sdm.cache import Cache
 from sdm.tensor.mixin import DeviceMixin
 
@@ -74,6 +75,21 @@ def __repr__(self) -> str:
         device_repr = f"device={device}" if device.type != "cpu" else ""
         return f"{self.__class__.__name__}({device_repr})"
 
+    # Helpers #################################################################
+
+    def _validate_feature_stype(self, x: TableTensor) -> None:
+        invalid = x.active_stypes - self.supported_feature_stypes - {Stype.id}
+        if len(invalid) > 0:
+            stypes = ", ".join(f"{str(stype)!r}" for stype in invalid)
+            warn_once(
+                key="model-unsupported-feature-stypes",
+                message=(
+                    f"{self.__class__.__name__!r} received unsupported "
+                    f"feature stypes {stypes}. Columns with unsupported "
+                    f"feature stypes will not be consumed by the model."
+                ),
+            )
+
 
 class _Prefetcher(Generic[T]):
     def __init__(

sdm/models/kumo/timeseries/anomaly/model.py

@@ -5,16 +5,19 @@
 
 from __future__ import annotations
 
+import copy
 from collections.abc import Sequence
 from typing import ClassVar
 
 import torch
 from torch import Tensor
 
 from sdm import Recipe, Stype, TableTensor
+from sdm._inference import inference_mode
 from sdm.models import Callback, Model
 from sdm.models._huggingface import download_checkpoint
 from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe
+from sdm.processing.execution import MemberQuery, RecipeExecution
 
 
 class KumoAnomaly(Model):
@@ -83,14 +86,14 @@ def default_recipe(cls) -> Recipe:
 
     def forward(
         self,
-        x: Tensor | TableTensor,
+        x: Tensor | TableTensor,  # [ ..., R, D]
         *,
         recipe: Recipe | None = None,
         num_samples: int = 30,
         num_steps: int = 20,
         callbacks: Sequence[Callback] | None = None,
         generator: torch.Generator | None = None,
-    ) -> TableTensor:
+    ) -> TableTensor:  # [..., R, 1]
         r"""The forward pass.
 
         Args:
@@ -111,7 +114,57 @@ def forward(
             absolute reconstruction error across preprocessed features at each
             time step.
         """
-        raise NotImplementedError
+        callbacks = () if callbacks is None else callbacks
+        requires_grad = self.training
+        requires_grad |= any(callback.requires_grad for callback in callbacks)
+
+        recipe_execution = RecipeExecution(
+            self.default_recipe() if recipe is None else copy.deepcopy(recipe)
+        )
+        with (
+            torch.amp.autocast(x.device.type, enabled=False),
+            inference_mode("no_grad" if requires_grad else "inference"),
+        ):
+            query = recipe_execution.fit_transform(
+                x=x,
+                y=None,
+                related_tables=None,
+                num_members=1,
+                generator=generator,
+            )[0]
+
+        for callback in callbacks:
+            query = MemberQuery(
+                *callback.on_query_preprocessing_end(
+                    model=self,
+                    x=query.x,
+                    related_tables=query.related_tables,
+                ),
+            )
+
+        self._validate_feature_stype(query.x)
+
+        with inference_mode("grad" if requires_grad else "inference"):
+            out = self.model(
+                x=query.x.numerical,
+                num_samples=num_samples,
+                num_steps=num_steps,
+                generator=generator,
+            )
+
+            out = TableTensor(
+                columns={Stype.numerical: ("mae",)},
+                numerical=out.to(query.x.numerical),
+            )
+
+            for callback in callbacks:
+                out = callback.on_model_forward_end(self, out)
+
+        with (
+            torch.amp.autocast(x.device.type, enabled=False),
+            inference_mode("grad" if requires_grad else "inference"),
+        ):
+            return recipe_execution.transform_output([out])[0]
 
 
 class _KumoAnomaly(torch.nn.Module):
@@ -121,3 +174,12 @@ def __init__(
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+
+    def forward(
+        self,
+        x: Tensor,  # [..., R, C]
+        num_samples: int,
+        num_steps: int,
+        generator: torch.Generator | None = None,
+    ) -> Tensor:  # [..., R, 1]
+        return x.mean(dim=-1, keepdim=True)