From 7b98067cf1340931a689c180ac483b14b86ca801 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Wed, 21 Oct 2026 10:48:05 +0000 Subject: [PATCH 01/15] update --- sdm/models/kumo/timeseries/anomaly/model.py | 100 ++++++++++++++++++- sdm/models/kumo/timeseries/forecast/model.py | 4 +- 2 files changed, 100 insertions(+), 4 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index e7cb377ac..8ff1a8886 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -3,10 +3,17 @@ # ruff: noqa: D205 +from __future__ import annotations + +from collections.abc import Sequence from typing import ClassVar -from sdm import Recipe, Stype -from sdm.models import Model +import torch +from torch import Tensor + +from sdm import Recipe, Stype, TableTensor +from sdm.models import Callback, Model +from sdm.models._huggingface import download_checkpoint class KumoAnomaly(Model): @@ -14,13 +21,102 @@ class KumoAnomaly(Model): `"NV-Tesseract-AD: Curriculum-Enhanced Conditional Diffusion for Multivariate Time-Series Anomaly Detection" `__ paper. + + :class:`KumoAnomaly` reconstructs multivariate time series with conditional + diffusion and uses reconstruction errors to score anomalies. It combines + sinusoidal time embeddings and learned feature embeddings with + diffusion-step and masking-strategy embeddings. Its denoiser alternates + attention over time and features with gated residuals and aggregated skip + connections that predict noise at each time step and feature. + + At inference, two complementary masks hide alternating temporal blocks. + Reverse diffusion reconstructs each hidden block from the visible context, + and the sample median combines repeated reconstructions. Mean absolute + reconstruction error across features gives a score per time step; larger + scores indicate less accurately reconstructed observations. + + Args: + pretrained: Whether to load the pretrained checkpoint. + device: The device. """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( {Stype.numerical} ) + def __init__( + self, + pretrained: bool = True, + device: torch.device | str | None = None, + ) -> None: + super().__init__() + + self.model = _KumoAnomaly( + device="meta" if pretrained else device, + ) + + if pretrained: + self._load_from_pretrained(device=device) + + self.eval() + + def _load_from_pretrained( + self, + device: torch.device | str | None, + ) -> KumoAnomaly: + device = torch.get_default_device() if device is None else device + + path = download_checkpoint( + repo_id="nvidia/Kumo-Anomaly", + filename="final_model.pth", + revision="226c5003b0582adc1cac433123f47324bfae39f2", + ) + torch.load(path, map_location=device, weights_only=True) # TODO + + return self + @classmethod def default_recipe(cls) -> Recipe: r""":meta private:""" # noqa: D415 return Recipe() + + def forward( + self, + x: Tensor | TableTensor, + *, + recipe: Recipe | None = None, + num_samples: int = 30, + num_steps: int = 20, + callbacks: Sequence[Callback] | None = None, + generator: torch.Generator | None = None, + ) -> TableTensor: + r"""The forward pass. + + Args: + x: The feature tensor with shape ``[..., R, D]`` with ``R`` rows + and ``D`` columns. + recipe: The custom recipe for pre- and post-processing. + num_samples: The number of reconstruction samples per complementary + masking strategy. Reconstructions are combined using the + sample median. + num_steps: The number of denoiser evaluations per sample and + masking strategy. + callbacks: Callbacks applied in sequence to this model call. + generator: Pseudorandom number generator used for sampling during + pre-processing and model execution. + + Returns: + The anomaly scores with shape ``[..., R, 1]``, denoting the mean + absolute reconstruction error across preprocessed features at each + time step. + """ + raise NotImplementedError + + +class _KumoAnomaly(torch.nn.Module): + def __init__( + self, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() diff --git a/sdm/models/kumo/timeseries/forecast/model.py b/sdm/models/kumo/timeseries/forecast/model.py index b312ef67c..7eb9e2568 100644 --- a/sdm/models/kumo/timeseries/forecast/model.py +++ b/sdm/models/kumo/timeseries/forecast/model.py @@ -94,14 +94,14 @@ def _load_from_pretrained( ) -> KumoForecast: device = torch.get_default_device() if device is None else device - checkpoint = download_checkpoint( + path = download_checkpoint( repo_id="nvidia/Kumo-Forecast", filename="run8_best_model_cr.pt" if use_cross_channel else "moment_head_512_6hr.pt", revision="abff20a58834638b28227ff4ab934f26206e4b09", ) - ckpt = torch.load(checkpoint, map_location=device, weights_only=True) + ckpt = torch.load(path, map_location=device, weights_only=True) self.model.load_state_dict(remap_ckpt(ckpt), assign=True) return self From 2aa400f3887a40cbce67d8b152cfa260385a90b9 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Wed, 21 Oct 2026 11:18:08 +0000 Subject: [PATCH 02/15] update --- sdm/models/kumo/timeseries/anomaly/model.py | 3 +- sdm/models/kumo/timeseries/anomaly/recipe.py | 15 +++++ sdm/processing/__init__.py | 2 + sdm/processing/numerical/__init__.py | 2 + sdm/processing/numerical/min_max_scale.py | 53 +++++++++++++++ sdm/processing/numerical/pca.py | 13 ++-- .../kumo/timeseries/anomaly/test_recipe.py | 45 +++++++++++++ .../numerical/test_min_max_scale.py | 67 +++++++++++++++++++ test/processing/numerical/test_pca.py | 18 +++++ test/processing/test_contract.py | 1 + test/processing/test_recipe.py | 4 +- 11 files changed, 215 insertions(+), 8 deletions(-) create mode 100644 sdm/models/kumo/timeseries/anomaly/recipe.py create mode 100644 sdm/processing/numerical/min_max_scale.py create mode 100644 test/models/kumo/timeseries/anomaly/test_recipe.py create mode 100644 test/processing/numerical/test_min_max_scale.py diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 8ff1a8886..46a3534a7 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -14,6 +14,7 @@ from sdm import Recipe, Stype, TableTensor from sdm.models import Callback, Model from sdm.models._huggingface import download_checkpoint +from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe class KumoAnomaly(Model): @@ -78,7 +79,7 @@ def _load_from_pretrained( @classmethod def default_recipe(cls) -> Recipe: r""":meta private:""" # noqa: D415 - return Recipe() + return default_recipe() def forward( self, diff --git a/sdm/models/kumo/timeseries/anomaly/recipe.py b/sdm/models/kumo/timeseries/anomaly/recipe.py new file mode 100644 index 000000000..3f6425b6c --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/recipe.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import sdm.processing as sp + + +def default_recipe() -> sp.Recipe: # noqa: D103 + return sp.Recipe( + features=sp.StypeDispatch( + numerical=[ + sp.PCA(num_components=40), + sp.MinMaxScale(), + ] + ), + ) diff --git a/sdm/processing/__init__.py b/sdm/processing/__init__.py index 45a6d4a48..800b423ab 100644 --- a/sdm/processing/__init__.py +++ b/sdm/processing/__init__.py @@ -35,6 +35,7 @@ RankGaussian, Standardize, RobustScale, + MinMaxScale, FlipSign, DropConstantColumns, PCA, @@ -83,6 +84,7 @@ "RankGaussian", "Standardize", "RobustScale", + "MinMaxScale", "FlipSign", "DropConstantColumns", "PCA", diff --git a/sdm/processing/numerical/__init__.py b/sdm/processing/numerical/__init__.py index d8665649d..673cfff32 100644 --- a/sdm/processing/numerical/__init__.py +++ b/sdm/processing/numerical/__init__.py @@ -13,6 +13,7 @@ from sdm.processing.numerical.rank_gaussian import RankGaussian from sdm.processing.numerical.standardize import Standardize from sdm.processing.numerical.robust_scale import RobustScale +from sdm.processing.numerical.min_max_scale import MinMaxScale from sdm.processing.numerical.flip_sign import FlipSign from sdm.processing.numerical.constant import DropConstantColumns from sdm.processing.numerical.pca import PCA @@ -32,6 +33,7 @@ "RankGaussian", "Standardize", "RobustScale", + "MinMaxScale", "FlipSign", "DropConstantColumns", "PCA", diff --git a/sdm/processing/numerical/min_max_scale.py b/sdm/processing/numerical/min_max_scale.py new file mode 100644 index 000000000..1b84258e3 --- /dev/null +++ b/sdm/processing/numerical/min_max_scale.py @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor + +from sdm import Stype, TableTensor +from sdm.processing import InvertibleMixin, Processor +from sdm.processing.numerical._stats import _isfinite + + +class MinMaxScale(Processor, InvertibleMixin): + """Scale each numerical column by its fitted minimum and range. + + Fitted minima map to zero and maxima to one. Values outside the fitted + range are not clipped. Constant columns use a unit scale. + NaN and infinite values are ignored when fitting statistics and preserved + during the transform. Columns without finite fitting values yield NaN. + """ + + handles_stypes = (Stype.numerical,) + requires_fit = True + + _minimum: Tensor + _scale: Tensor + + def __init__(self) -> None: + super().__init__() + self.register_buffer("_minimum", torch.empty(0)) + self.register_buffer("_scale", torch.empty(0)) + + def _fit( + self, + table: TableTensor, + *, + generator: torch.Generator | None = None, + ) -> None: + x = table.numerical + missing = ~_isfinite(x) + minimum = x.masked_fill(missing, torch.inf).amin(-2, keepdim=True) + maximum = x.masked_fill(missing, -torch.inf).amax(-2, keepdim=True) + minimum.masked_fill_(minimum.isinf(), torch.nan) + scale = maximum - minimum + self._set_buffer("_minimum", minimum) + self._set_buffer("_scale", scale.masked_fill_(scale == 0, 1)) + + def _transform(self, table: TableTensor) -> TableTensor: + numerical = (table.numerical - self._minimum).div_(self._scale) + return table.replace_blocks(numerical=numerical) + + def _inverse_transform(self, table: TableTensor) -> TableTensor: + numerical = (table.numerical * self._scale).add_(self._minimum) + return table.replace_blocks(numerical=numerical) diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index 2924ec561..c993c3b75 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -14,7 +14,8 @@ class PCA(Processor): r"""Project numerical columns onto their principal components. Args: - num_components: Number of principal components to keep. + num_components: Number of principal components to keep. Numerical + columns are unchanged when their count is at most this value. """ handles_stypes = (Stype.numerical,) @@ -36,11 +37,8 @@ def _fit( generator: torch.Generator | None = None, ) -> None: - if table.numerical.size(-1) == 0: - raise ValueError( - f"{self.__class__.__name__!r} requires 'table' to have " - f"numerical features" - ) + if table.numerical.size(-1) <= self.num_components: + return self._set_buffer("mean", table.numerical.mean(dim=-2, keepdim=True)) _, _, vh = torch.linalg.svd( @@ -58,6 +56,9 @@ def _fit( ) def _transform(self, table: TableTensor) -> TableTensor: + if table.numerical.size(-1) <= self.num_components: + return table + x = (table.numerical - self.mean) @ self.components out = TableTensor( columns={Stype.numerical: [f"pca_{i}" for i in range(x.size(-1))]}, diff --git a/test/models/kumo/timeseries/anomaly/test_recipe.py b/test/models/kumo/timeseries/anomaly/test_recipe.py new file mode 100644 index 000000000..622ac87e2 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_recipe.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm import EnsembleTable, TableTensor +from sdm.models.kumo.timeseries.anomaly import KumoAnomaly +from sdm.processing import PCA + + +@pytest.mark.parametrize("num_features", [3, 40, 45]) +@pytest.mark.parametrize("batch_shape", [(), (2,)]) +def test_recipe_reduces_and_scales_features( + num_features: int, + batch_shape: tuple[int, ...], +) -> None: + reference = TableTensor.from_tensor( + torch.randn(*batch_shape, 50, num_features) + ) + query = TableTensor.from_tensor( + torch.randn(*batch_shape, 10, num_features) + ) + recipe = KumoAnomaly.default_recipe() + recipe.features.fit_ensemble( + EnsembleTable.from_table(reference, num_members=1) + ) + output = recipe.features.transform_ensemble( + EnsembleTable.from_table(query, num_members=1) + )[0] + + fitted = reference + expected = query + if num_features > 40: + pca = PCA(num_components=40).fit(reference) + fitted = pca.transform(reference) + expected = pca.transform(query) + + minimum = fitted.numerical.amin(dim=-2, keepdim=True) + scale = fitted.numerical.amax(dim=-2, keepdim=True) - minimum + scale = scale.masked_fill(scale == 0, 1) + assert output.columns == expected.columns + torch.testing.assert_close( + output.numerical, (expected.numerical - minimum) / scale + ) diff --git a/test/processing/numerical/test_min_max_scale.py b/test/processing/numerical/test_min_max_scale.py new file mode 100644 index 000000000..ed83c37e4 --- /dev/null +++ b/test/processing/numerical/test_min_max_scale.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm import TableTensor +from sdm.processing import MinMaxScale +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize( + "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64] +) +def test_min_max_scale_fitted_range( + device: torch.device, dtype: torch.dtype +) -> None: + reference = torch.tensor( + [[1.0, 5.0], [3.0, 5.0]], device=device, dtype=dtype + ) + processor = MinMaxScale() + out = processor.fit_transform(TableTensor.from_tensor(reference)) + torch.testing.assert_close( + out.numerical, reference.new_tensor([[0.0, 0.0], [1.0, 0.0]]) + ) + + query = reference.new_tensor([[-1.0, 3.0], [5.0, 7.0]]) + out = processor.transform(TableTensor.from_tensor(query)) + torch.testing.assert_close( + out.numerical, query.new_tensor([[-1.0, -2.0], [2.0, 2.0]]) + ) + torch.testing.assert_close( + processor.inverse_transform(out).numerical, query + ) + + +@withCUDA +def test_min_max_scale_batched_statistics(device: torch.device) -> None: + reference = torch.tensor( + [[[[1.0], [3.0]], [[10.0], [18.0]]]], device=device + ) + query = torch.tensor([[[[2.0]], [[16.0]]]], device=device) + processor = MinMaxScale().fit(TableTensor.from_tensor(reference)) + out = processor.transform(TableTensor.from_tensor(query)) + torch.testing.assert_close( + out.numerical, query.new_tensor([[[[0.5]], [[0.75]]]]) + ) + + +@withCUDA +def test_min_max_scale_nonfinite_values(device: torch.device) -> None: + x = torch.tensor( + [ + [1.0, torch.nan], + [3.0, torch.inf], + [torch.inf, -torch.inf], + [-torch.inf, torch.nan], + [torch.nan, torch.nan], + ], + device=device, + ) + out = MinMaxScale().fit_transform(TableTensor.from_tensor(x)) + expected = x.clone() + expected[:2, 0] = x.new_tensor([0.0, 1.0]) + expected[:, 1] = torch.nan + torch.testing.assert_close(out.numerical, expected, equal_nan=True) diff --git a/test/processing/numerical/test_pca.py b/test/processing/numerical/test_pca.py index 964c72728..049c5432d 100644 --- a/test/processing/numerical/test_pca.py +++ b/test/processing/numerical/test_pca.py @@ -19,6 +19,24 @@ def test_basic(shape: tuple[int, ...]) -> None: assert output.numerical_columns == ("pca_0", "pca_1") +@pytest.mark.parametrize("num_components", [3, 5]) +@pytest.mark.parametrize("shape", [(2, 3), (2, 4, 3)]) +def test_pca_preserves_features_without_reduction( + num_components: int, + shape: tuple[int, ...], +) -> None: + table = TableTensor.from_tensor(torch.randn(shape)) + pca = PCA(num_components=num_components) + output = pca.fit_transform(table) + assert output.columns == table.columns + torch.testing.assert_close(output.numerical, table.numerical) + + query = TableTensor.from_tensor(torch.randn(shape)) + output = pca.transform(query) + assert output.columns == query.columns + torch.testing.assert_close(output.numerical, query.numerical) + + def test_pca_projects_onto_fitted_dominant_direction() -> None: steps = torch.arange(10, dtype=torch.float) fit_data = torch.stack((steps, steps), dim=-1) diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index 7525b3b8c..8078d86d6 100644 --- a/test/processing/test_contract.py +++ b/test/processing/test_contract.py @@ -129,6 +129,7 @@ def _make_processor_pair( ProcessorCase(sp.Standardize()), ProcessorCase(sp.RankGaussian()), ProcessorCase(sp.RobustScale()), + ProcessorCase(sp.MinMaxScale()), ProcessorCase(sp.FlipSign()), ProcessorCase(sp.DropConstantColumns()), ProcessorCase(sp.PCA(2)), diff --git a/test/processing/test_recipe.py b/test/processing/test_recipe.py index faa05f0de..42b2eb973 100644 --- a/test/processing/test_recipe.py +++ b/test/processing/test_recipe.py @@ -24,7 +24,9 @@ def test_recipe_normalizes_empty_roles_and_repr() -> None: assert len(recipe.features) == 1 assert len(recipe.output) == 0 assert "features=Sequential" in repr(recipe) - assert "target=Identity()" in repr(recipe) + assert "target=" not in repr(recipe) + assert "target=" not in repr(sp.Recipe(target=sp.Identity())) + assert "target=Standardize" in repr(sp.Recipe(target=sp.Standardize())) def test_target_forward_then_inverse_round_trips() -> None: From 3524ba79e2f828f9c17043e4d2a4f1299dd5f1ac Mon Sep 17 00:00:00 2001 From: rusty1s Date: Wed, 21 Oct 2026 11:22:06 +0000 Subject: [PATCH 03/15] update --- sdm/processing/numerical/pca.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index c993c3b75..c1e37e9ee 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -14,8 +14,7 @@ class PCA(Processor): r"""Project numerical columns onto their principal components. Args: - num_components: Number of principal components to keep. Numerical - columns are unchanged when their count is at most this value. + num_components: Number of principal components to keep. """ handles_stypes = (Stype.numerical,) @@ -45,11 +44,7 @@ def _fit( table.numerical - self.mean, full_matrices=False, ) - num_components = min( - self.num_components, - table.numerical.size(-2), - table.numerical.size(-1), - ) + num_components = min(self.num_components, table.numerical.size(-2)) self._set_buffer( name="components", value=vh[..., :num_components, :].transpose(-2, -1), From 975d63a03e7c7276ee0170a34df1dfed7773d44c Mon Sep 17 00:00:00 2001 From: rusty1s Date: Wed, 21 Oct 2026 11:25:16 +0000 Subject: [PATCH 04/15] update --- test/processing/numerical/test_pca.py | 18 ------------------ test/processing/test_recipe.py | 4 +--- 2 files changed, 1 insertion(+), 21 deletions(-) diff --git a/test/processing/numerical/test_pca.py b/test/processing/numerical/test_pca.py index 049c5432d..964c72728 100644 --- a/test/processing/numerical/test_pca.py +++ b/test/processing/numerical/test_pca.py @@ -19,24 +19,6 @@ def test_basic(shape: tuple[int, ...]) -> None: assert output.numerical_columns == ("pca_0", "pca_1") -@pytest.mark.parametrize("num_components", [3, 5]) -@pytest.mark.parametrize("shape", [(2, 3), (2, 4, 3)]) -def test_pca_preserves_features_without_reduction( - num_components: int, - shape: tuple[int, ...], -) -> None: - table = TableTensor.from_tensor(torch.randn(shape)) - pca = PCA(num_components=num_components) - output = pca.fit_transform(table) - assert output.columns == table.columns - torch.testing.assert_close(output.numerical, table.numerical) - - query = TableTensor.from_tensor(torch.randn(shape)) - output = pca.transform(query) - assert output.columns == query.columns - torch.testing.assert_close(output.numerical, query.numerical) - - def test_pca_projects_onto_fitted_dominant_direction() -> None: steps = torch.arange(10, dtype=torch.float) fit_data = torch.stack((steps, steps), dim=-1) diff --git a/test/processing/test_recipe.py b/test/processing/test_recipe.py index 42b2eb973..faa05f0de 100644 --- a/test/processing/test_recipe.py +++ b/test/processing/test_recipe.py @@ -24,9 +24,7 @@ def test_recipe_normalizes_empty_roles_and_repr() -> None: assert len(recipe.features) == 1 assert len(recipe.output) == 0 assert "features=Sequential" in repr(recipe) - assert "target=" not in repr(recipe) - assert "target=" not in repr(sp.Recipe(target=sp.Identity())) - assert "target=Standardize" in repr(sp.Recipe(target=sp.Standardize())) + assert "target=Identity()" in repr(recipe) def test_target_forward_then_inverse_round_trips() -> None: From 47679cff77472c0b2d120985ab81bd99aa30df97 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Wed, 21 Oct 2026 11:33:09 +0000 Subject: [PATCH 05/15] update --- sdm/processing/numerical/min_max_scale.py | 23 ++++----- .../kumo/timeseries/anomaly/test_recipe.py | 45 ----------------- .../numerical/test_min_max_scale.py | 50 ++----------------- 3 files changed, 16 insertions(+), 102 deletions(-) delete mode 100644 test/models/kumo/timeseries/anomaly/test_recipe.py diff --git a/sdm/processing/numerical/min_max_scale.py b/sdm/processing/numerical/min_max_scale.py index 1b84258e3..4b68a78d4 100644 --- a/sdm/processing/numerical/min_max_scale.py +++ b/sdm/processing/numerical/min_max_scale.py @@ -12,22 +12,20 @@ class MinMaxScale(Processor, InvertibleMixin): """Scale each numerical column by its fitted minimum and range. - Fitted minima map to zero and maxima to one. Values outside the fitted - range are not clipped. Constant columns use a unit scale. NaN and infinite values are ignored when fitting statistics and preserved - during the transform. Columns without finite fitting values yield NaN. + during the transform. """ handles_stypes = (Stype.numerical,) requires_fit = True - _minimum: Tensor - _scale: Tensor + minimum: Tensor + scale: Tensor def __init__(self) -> None: super().__init__() - self.register_buffer("_minimum", torch.empty(0)) - self.register_buffer("_scale", torch.empty(0)) + self.register_buffer("minimum", torch.empty(0)) + self.register_buffer("scale", torch.empty(0)) def _fit( self, @@ -38,16 +36,17 @@ def _fit( x = table.numerical missing = ~_isfinite(x) minimum = x.masked_fill(missing, torch.inf).amin(-2, keepdim=True) + minimum[minimum.isinf()] = torch.nan maximum = x.masked_fill(missing, -torch.inf).amax(-2, keepdim=True) - minimum.masked_fill_(minimum.isinf(), torch.nan) scale = maximum - minimum - self._set_buffer("_minimum", minimum) - self._set_buffer("_scale", scale.masked_fill_(scale == 0, 1)) + scale[scale == 0] = 1 + self._set_buffer("minimum", minimum) + self._set_buffer("scale", scale) def _transform(self, table: TableTensor) -> TableTensor: - numerical = (table.numerical - self._minimum).div_(self._scale) + numerical = (table.numerical - self.minimum).div_(self.scale) return table.replace_blocks(numerical=numerical) def _inverse_transform(self, table: TableTensor) -> TableTensor: - numerical = (table.numerical * self._scale).add_(self._minimum) + numerical = (table.numerical * self.scale).add_(self.minimum) return table.replace_blocks(numerical=numerical) diff --git a/test/models/kumo/timeseries/anomaly/test_recipe.py b/test/models/kumo/timeseries/anomaly/test_recipe.py deleted file mode 100644 index 622ac87e2..000000000 --- a/test/models/kumo/timeseries/anomaly/test_recipe.py +++ /dev/null @@ -1,45 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import pytest -import torch - -from sdm import EnsembleTable, TableTensor -from sdm.models.kumo.timeseries.anomaly import KumoAnomaly -from sdm.processing import PCA - - -@pytest.mark.parametrize("num_features", [3, 40, 45]) -@pytest.mark.parametrize("batch_shape", [(), (2,)]) -def test_recipe_reduces_and_scales_features( - num_features: int, - batch_shape: tuple[int, ...], -) -> None: - reference = TableTensor.from_tensor( - torch.randn(*batch_shape, 50, num_features) - ) - query = TableTensor.from_tensor( - torch.randn(*batch_shape, 10, num_features) - ) - recipe = KumoAnomaly.default_recipe() - recipe.features.fit_ensemble( - EnsembleTable.from_table(reference, num_members=1) - ) - output = recipe.features.transform_ensemble( - EnsembleTable.from_table(query, num_members=1) - )[0] - - fitted = reference - expected = query - if num_features > 40: - pca = PCA(num_components=40).fit(reference) - fitted = pca.transform(reference) - expected = pca.transform(query) - - minimum = fitted.numerical.amin(dim=-2, keepdim=True) - scale = fitted.numerical.amax(dim=-2, keepdim=True) - minimum - scale = scale.masked_fill(scale == 0, 1) - assert output.columns == expected.columns - torch.testing.assert_close( - output.numerical, (expected.numerical - minimum) / scale - ) diff --git a/test/processing/numerical/test_min_max_scale.py b/test/processing/numerical/test_min_max_scale.py index ed83c37e4..984b784af 100644 --- a/test/processing/numerical/test_min_max_scale.py +++ b/test/processing/numerical/test_min_max_scale.py @@ -1,7 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest import torch from sdm import TableTensor @@ -10,22 +9,15 @@ @withCUDA -@pytest.mark.parametrize( - "dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64] -) -def test_min_max_scale_fitted_range( - device: torch.device, dtype: torch.dtype -) -> None: - reference = torch.tensor( - [[1.0, 5.0], [3.0, 5.0]], device=device, dtype=dtype - ) +def test_min_max_scale(device: torch.device) -> None: + data = torch.tensor([[1.0, 5.0], [3.0, 5.0]], device=device) processor = MinMaxScale() - out = processor.fit_transform(TableTensor.from_tensor(reference)) + out = processor.fit_transform(TableTensor.from_tensor(data)) torch.testing.assert_close( - out.numerical, reference.new_tensor([[0.0, 0.0], [1.0, 0.0]]) + out.numerical, data.new_tensor([[0.0, 0.0], [1.0, 0.0]]) ) - query = reference.new_tensor([[-1.0, 3.0], [5.0, 7.0]]) + query = data.new_tensor([[-1.0, 3.0], [5.0, 7.0]]) out = processor.transform(TableTensor.from_tensor(query)) torch.testing.assert_close( out.numerical, query.new_tensor([[-1.0, -2.0], [2.0, 2.0]]) @@ -33,35 +25,3 @@ def test_min_max_scale_fitted_range( torch.testing.assert_close( processor.inverse_transform(out).numerical, query ) - - -@withCUDA -def test_min_max_scale_batched_statistics(device: torch.device) -> None: - reference = torch.tensor( - [[[[1.0], [3.0]], [[10.0], [18.0]]]], device=device - ) - query = torch.tensor([[[[2.0]], [[16.0]]]], device=device) - processor = MinMaxScale().fit(TableTensor.from_tensor(reference)) - out = processor.transform(TableTensor.from_tensor(query)) - torch.testing.assert_close( - out.numerical, query.new_tensor([[[[0.5]], [[0.75]]]]) - ) - - -@withCUDA -def test_min_max_scale_nonfinite_values(device: torch.device) -> None: - x = torch.tensor( - [ - [1.0, torch.nan], - [3.0, torch.inf], - [torch.inf, -torch.inf], - [-torch.inf, torch.nan], - [torch.nan, torch.nan], - ], - device=device, - ) - out = MinMaxScale().fit_transform(TableTensor.from_tensor(x)) - expected = x.clone() - expected[:2, 0] = x.new_tensor([0.0, 1.0]) - expected[:, 1] = torch.nan - torch.testing.assert_close(out.numerical, expected, equal_nan=True) From e7233bc10e4a4ade4386b3dac1eb93b352846d33 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:43:34 -0700 Subject: [PATCH 06/15] Add `KumoAnomaly` reconstruction and scoring Co-authored-by: Aditi Gautam --- sdm/models/kumo/timeseries/anomaly/model.py | 189 +++++++++++++++++- .../kumo/timeseries/anomaly/test_model.py | 88 ++++++++ 2 files changed, 274 insertions(+), 3 deletions(-) create mode 100644 test/models/kumo/timeseries/anomaly/test_model.py diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 46a3534a7..51178bc97 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -5,16 +5,25 @@ from __future__ import annotations +import copy from collections.abc import Sequence -from typing import ClassVar +from functools import partial +from typing import Any, ClassVar, cast import torch +import torch.nn.functional as F from torch import Tensor from sdm import Recipe, Stype, TableTensor +from sdm._inference import inference_mode +from sdm._warnings import warn_once from sdm.models import Callback, Model from sdm.models._huggingface import download_checkpoint +from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser +from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule class KumoAnomaly(Model): @@ -96,7 +105,7 @@ def forward( Args: x: The feature tensor with shape ``[..., R, D]`` with ``R`` rows and ``D`` columns. - recipe: The custom recipe for pre- and post-processing. + recipe: The custom recipe for pre-processing. num_samples: The number of reconstruction samples per complementary masking strategy. Reconstructions are combined using the sample median. @@ -111,13 +120,187 @@ 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) + + param = next(self.parameters()) + if not isinstance(x, TableTensor): + x = TableTensor.from_tensor(x) + x = cast(TableTensor, x.to(param.device)) + + recipe = ( + self.default_recipe() if recipe is None else copy.deepcopy(recipe) + ) + with ( + torch.amp.autocast(param.device.type, enabled=False), + inference_mode("no_grad" if requires_grad else "inference"), + ): + x = recipe.features.fit_transform(x, generator=generator) + + for callback in callbacks: + x, _ = callback.on_query_preprocessing_end(self, x, 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." + ), + ) + if x.numerical.size(-1) > self.model.num_features: + raise ValueError( + f"{self.__class__.__name__!r} supports at most " + f"{self.model.num_features} features after pre-processing " + f"(got {x.numerical.size(-1)})" + ) + + with inference_mode("grad" if requires_grad else "inference"): + out = self.model( + x.numerical.to(param.dtype), + num_samples=num_samples, + num_steps=num_steps, + generator=generator, + ) + out = TableTensor( + columns={Stype.numerical: ["score"]}, + numerical=out.unsqueeze(-1), + ) + for callback in callbacks: + out = callback.on_model_forward_end(self, out) + return cast(TableTensor, out.to(x.dtype)) class _KumoAnomaly(torch.nn.Module): def __init__( self, + num_features: int = 40, + time_channels: int = 256, + feature_channels: int = 64, + channels: int = 128, + num_layers: int = 6, + num_heads: int = 8, + embedding_channels: int = 256, + num_steps: int = 500, + window_size: int = 100, + num_blocks: int = 4, + batch_size: int = 32, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + + self.num_features = num_features + self.num_steps = num_steps + self.window_size = window_size + self.num_blocks = num_blocks + self.batch_size = batch_size + + self.side_embedding = SideEmbedding( + num_features=num_features, + time_channels=time_channels, + feature_channels=feature_channels, + **factory_kwargs, + ) + self.denoiser = DiffusionDenoiser( + channels=channels, + side_channels=time_channels + feature_channels, + num_steps=num_steps, + num_layers=num_layers, + num_heads=num_heads, + embedding_channels=embedding_channels, + **factory_kwargs, + ) + + def forward( + self, + x: Tensor, # [..., R, D] + num_samples: int, + num_steps: int, + generator: torch.Generator | None, + ) -> Tensor: # [..., R] + + *size, length, features = x.size() + x = x.reshape(-1, length, features) + x = F.pad(x, (0, self.num_features - features)) + + # Pad the last window by repeating the last row: + padding = -length % self.window_size + x = torch.cat([x, x[:, -1:].expand(-1, padding, -1)], dim=1) + windows = x.reshape(-1, self.window_size, self.num_features) + + schedule = DiffusionSchedule( + num_steps=self.num_steps, + beta_start=0.0001, + beta_end=0.01, + schedule="linear", + device="cpu", + ) + sampler = DPMSolver(schedule.beta.float(), num_steps).to(x.device) + + out = torch.cat( + [ + self.reconstruct(batch, sampler, num_samples, generator) + for batch in windows.split(self.batch_size) + ] + ) + out = out.reshape(x.size())[:, :length, :features] + out = (out - x[:, :length, :features]).abs().mean(dim=-1) + return out.reshape(*size, length) + + def reconstruct( + self, + x: Tensor, # [B, L, K] + sampler: DPMSolver, + num_samples: int, + generator: torch.Generator | None, + ) -> Tensor: # [B, L, K] + + x = x.transpose(1, 2) # [B, K, L] + batch, _, length = x.size() + mean = x.mean(dim=(1, 2), keepdim=True) + std = x.std(dim=(1, 2), keepdim=True) + 1e-5 + x = ((x - mean) / std).clamp(-10, 10) + + time = torch.arange(length, device=x.device).expand(batch, -1) + blocks = time // (length // self.num_blocks) + + out = x.new_zeros((num_samples, *x.size())) + for strategy in range(2): + mask = ((blocks % 2) == strategy).to(x.dtype) + mask = mask.unsqueeze(1).expand_as(x) # [B, K, L] + predict_noise = partial( + self.predict_noise, + cond=mask * x, + target_mask=1 - mask, + side_info=self.side_embedding(time, mask), + strategy=torch.tensor([strategy], device=x.device), + ) + for i in range(num_samples): + noise = torch.empty_like(x).normal_(generator=generator) + sample = sampler(noise, predict_noise) * std + mean + out[i] += (1 - mask) * sample + + return out.median(dim=0).values.transpose(1, 2) + + def predict_noise( + self, + x: Tensor, # [B, K, L] + step: Tensor, # [1] + *, + cond: Tensor, # [B, K, L] + target_mask: Tensor, # [B, K, L] + side_info: Tensor, # [B, S, K, L] + strategy: Tensor, # [1] + ) -> Tensor: # [B, K, L] + return self.denoiser( + x=torch.stack([cond, target_mask * x], dim=1), + side_info=side_info, + step=step, + strategy=strategy, + ) diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py new file mode 100644 index 000000000..04a46f0ad --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from sdm import CategoricalTensor, Recipe, Stype, TableTensor +from sdm.explain.gradient import CaptureInputs, EnableInputGradients +from sdm.models import Callback, KumoAnomaly +from sdm.testing import withCUDA + + +@withCUDA +def test_forward(device: torch.device) -> None: + model = KumoAnomaly(pretrained=False, device=device) + x = torch.randn(2, 130, 3, device=device) + + def generator() -> torch.Generator: + return torch.Generator(device).manual_seed(0) + + out = model(x, num_samples=1, num_steps=1, generator=generator()) + assert out.size() == (2, 130, 1) + assert out.columns["numerical"] == ("score",) + assert out.device == device + assert out.numerical.isfinite().all() + assert (out.numerical >= 0).all() + + again = model(x, num_samples=1, num_steps=1, generator=generator()) + assert torch.equal(again.numerical, out.numerical) + + +def test_wide_float64_forward() -> None: + model = KumoAnomaly(pretrained=False) + x = torch.randn(60, 50, dtype=torch.float64) + + out = model(x, num_samples=1, num_steps=1) + assert out.size() == (60, 1) + assert out.numerical.dtype == torch.float64 + + +class DoubleScores(Callback): + def on_model_forward_end( + self, + model: torch.nn.Module, + out: TableTensor, + ) -> TableTensor: + return out.replace_blocks(numerical=2 * out.numerical.double()) + + +def test_callbacks() -> None: + model = KumoAnomaly(pretrained=False) + x = torch.randn(30, 2) + capture = CaptureInputs() + callbacks = [EnableInputGradients(), capture, DoubleScores()] + + def generator() -> torch.Generator: + return torch.Generator().manual_seed(0) + + out = model(x, num_samples=1, num_steps=1, generator=generator()) + doubled = model( + x, + num_samples=1, + num_steps=1, + callbacks=callbacks, + generator=generator(), + ) + assert doubled.numerical.dtype == torch.float32 + torch.testing.assert_close(doubled.numerical, 2 * out.numerical) + + doubled.numerical.sum().backward() + ((x, _),) = capture.inputs + assert x.numerical.grad is not None + + +def test_invalid_inputs() -> None: + model = KumoAnomaly(pretrained=False) + x = TableTensor( + columns={Stype.numerical: ["a"], Stype.categorical: ["b"]}, + numerical=torch.randn(10, 1), + categorical=CategoricalTensor.from_tensor( + torch.randint(0, 5, (10, 1)) + ), + ) + + with pytest.warns(RuntimeWarning, match="unsupported feature stypes"): + model(x, num_samples=1, num_steps=1) + with pytest.raises(ValueError, match="at most 40 features"): + model(torch.randn(10, 41), recipe=Recipe()) From cbc7404d92c2f445e6857cb26d3edd680c41fea1 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 03:14:30 +0000 Subject: [PATCH 07/15] update --- sdm/models/kumo/timeseries/anomaly/model.py | 73 ++++++++++++++++++- sdm/processing/numerical/min_max_scale.py | 2 +- .../kumo/timeseries/anomaly/test_model.py | 20 +++++ 3 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 test/models/kumo/timeseries/anomaly/test_model.py diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 46a3534a7..bc749aef7 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -5,6 +5,7 @@ from __future__ import annotations +import copy from collections.abc import Sequence from typing import ClassVar @@ -12,9 +13,11 @@ 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 MemberContext, 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,62 @@ 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"), + ): + context = recipe_execution.fit_transform( + x=x, + y=None, + related_tables=None, + num_members=1, + generator=generator, + )[0] + + for callback in callbacks: + context = MemberContext( + *callback.on_context_preprocessing_end( + model=self, + x=context.x, + y=context.y, + x_context_only=context.x_context_only, + related_tables=context.related_tables, + ), + input_stypes=context.input_stypes, + ) + + self._validate_context( + x=context.x, + y=context.y, + x_context_only=context.x_context_only, + related_tables=context.related_tables, + ) + + with inference_mode("grad" if requires_grad else "inference"): + out = self.model( + context.x.numerical, + num_samples=num_samples, + num_steps=num_steps, + generator=generator, + ) + + for callback in callbacks: + out = callback.on_model_forward_end(self, out) + + out = out.to(context.x.numerical) + + 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 +179,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) diff --git a/sdm/processing/numerical/min_max_scale.py b/sdm/processing/numerical/min_max_scale.py index 4b68a78d4..191509345 100644 --- a/sdm/processing/numerical/min_max_scale.py +++ b/sdm/processing/numerical/min_max_scale.py @@ -10,7 +10,7 @@ class MinMaxScale(Processor, InvertibleMixin): - """Scale each numerical column by its fitted minimum and range. + """Scale each numerical column by its fitted min-max range. NaN and infinite values are ignored when fitting statistics and preserved during the transform. diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py new file mode 100644 index 000000000..996237d59 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch + +from sdm import TableTensor +from sdm.models import KumoAnomaly +from sdm.testing import withCUDA + + +@withCUDA +def test_forward(device: torch.device) -> None: + model = KumoAnomaly(pretrained=False, device=device) + x = TableTensor.from_tensor(torch.randn(32, 5, device=device)) + + out = model(x) + assert out.size() == (32, 1) + assert out.device == device + assert out.numerical.isfinite().all() + assert "mae" in out.numerical_columns From a0ad470af15c194ff83832a660e08128a5d282e6 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 08:34:42 +0000 Subject: [PATCH 08/15] update --- sdm/models/base/icl.py | 35 +---- sdm/models/base/model.py | 18 ++- sdm/models/kumo/timeseries/anomaly/model.py | 29 ++-- sdm/processing/execution.py | 165 +++++++++++--------- 4 files changed, 124 insertions(+), 123 deletions(-) diff --git a/sdm/models/base/icl.py b/sdm/models/base/icl.py index 560db6a21..78734c889 100644 --- a/sdm/models/base/icl.py +++ b/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, diff --git a/sdm/models/base/model.py b/sdm/models/base/model.py index a266c9960..1b48d1ecf 100644 --- a/sdm/models/base/model.py +++ b/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__( diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index bc749aef7..20a094b27 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -17,7 +17,7 @@ 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 MemberContext, RecipeExecution +from sdm.processing.execution import MemberQuery, RecipeExecution class KumoAnomaly(Model): @@ -125,7 +125,7 @@ def forward( torch.amp.autocast(x.device.type, enabled=False), inference_mode("no_grad" if requires_grad else "inference"), ): - context = recipe_execution.fit_transform( + query = recipe_execution.fit_transform( x=x, y=None, related_tables=None, @@ -134,27 +134,19 @@ def forward( )[0] for callback in callbacks: - context = MemberContext( - *callback.on_context_preprocessing_end( + query = MemberQuery( + *callback.on_query_preprocessing_end( model=self, - x=context.x, - y=context.y, - x_context_only=context.x_context_only, - related_tables=context.related_tables, + x=query.x, + related_tables=query.related_tables, ), - input_stypes=context.input_stypes, ) - self._validate_context( - x=context.x, - y=context.y, - x_context_only=context.x_context_only, - related_tables=context.related_tables, - ) + self._validate_feature_stype(query.x) with inference_mode("grad" if requires_grad else "inference"): out = self.model( - context.x.numerical, + x=query.x.numerical, num_samples=num_samples, num_steps=num_steps, generator=generator, @@ -163,7 +155,10 @@ def forward( for callback in callbacks: out = callback.on_model_forward_end(self, out) - out = out.to(context.x.numerical) + out = TableTensor( + columns={Stype.numerical: ("mae",)}, + numerical=out.to(query.x.numerical), + ) with ( torch.amp.autocast(x.device.type, enabled=False), diff --git a/sdm/processing/execution.py b/sdm/processing/execution.py index 16c61f799..68a4b0e08 100644 --- a/sdm/processing/execution.py +++ b/sdm/processing/execution.py @@ -5,7 +5,7 @@ import copy from collections.abc import Mapping, Sequence -from typing import NamedTuple, cast +from typing import NamedTuple, cast, overload import torch from torch import Tensor @@ -53,6 +53,7 @@ def num_members(self) -> int: assert self._y_locations is not None return len(self._y_locations) + @overload def fit_transform( self, x: Tensor | TableTensor, @@ -62,60 +63,89 @@ def fit_transform( num_members: int | None = None, context_only_columns: Sequence[str] | Sequence[int] | None = None, generator: torch.Generator | None = None, - ) -> tuple[MemberContext, ...]: + ) -> tuple[MemberContext, ...]: ... + + @overload + def fit_transform( + self, + x: Tensor | TableTensor, + y: None, + related_tables: RelatedTables | None, + *, + num_members: int | None = None, + context_only_columns: None = None, + generator: torch.Generator | None = None, + ) -> tuple[MemberQuery, ...]: ... + + def fit_transform( + self, + x: Tensor | TableTensor, + y: Tensor | TableTensor | None, + related_tables: RelatedTables | None, + *, + num_members: int | None = None, + context_only_columns: Sequence[str] | Sequence[int] | None = None, + generator: torch.Generator | None = None, + ) -> tuple[MemberContext | MemberQuery, ...]: """Fit and transform context data.""" + self._num_estimators = num_members + # Transform target first to be able to resolve task type. Inverse # target transforms require distinct member assignment: - y: EnsembleTable = _to_ensemble_table(y, num_members, expand=True) - - series_processors = tuple( - module - for processor in (self.recipe.features, self.recipe.target) - for module in _iter_modules(processor) - if isinstance(module, sp.SeriesEnsembleProcessor) - ) - for module in series_processors: - module._set_buffer( - name="_offset", - value=torch.zeros( - size=(len(y),), - dtype=torch.long, - device=y[0].device, - ), + series_processors: tuple[sp.SeriesEnsembleProcessor, ...] = () + if y is not None: + y: EnsembleTable = _to_ensemble_table(y, num_members, expand=True) + + series_processors = tuple( + module + for processor in (self.recipe.features, self.recipe.target) + for module in _iter_modules(processor) + if isinstance(module, sp.SeriesEnsembleProcessor) ) + for module in series_processors: + module._set_buffer( + name="_offset", + value=torch.zeros( + size=(len(y),), + dtype=torch.long, + device=y[0].device, + ), + ) - y = self.recipe.target.fit_transform_ensemble(y, generator=generator) - self._num_estimators = num_members - self._y_locations = y._locations + y = self.recipe.target.fit_transform_ensemble( + y, generator=generator + ) + self._y_locations = y._locations - task_dispatchers = tuple( - module - for processor in (self.recipe.features, self.recipe.output) - for module in _iter_modules(processor) - if isinstance(module, sp.TaskDispatch) - ) - if len(task_dispatchers) > 0: - if any(group.size(-1) != 1 for group in y._iter_groups()): - raise ValueError( - "Expected the transformed target to contain exactly one " - "column" - ) + task_dispatchers = tuple( + module + for processor in (self.recipe.features, self.recipe.output) + for module in _iter_modules(processor) + if isinstance(module, sp.TaskDispatch) + ) + if len(task_dispatchers) > 0: + if any(group.size(-1) != 1 for group in y._iter_groups()): + raise ValueError( + "Expected the transformed target to contain exactly " + "one column" + ) - if all( - group.numerical.size(-1) == 1 for group in y._iter_groups() - ): - task = "regression" - elif all( - group.categorical.size(-1) == 1 for group in y._iter_groups() - ): - task = "classification" - else: - raise ValueError( - "'Recipe.target' must resolve to a single task type" - ) + if all( + group.numerical.size(-1) == 1 for group in y._iter_groups() + ): + task = "regression" + elif all( + group.categorical.size(-1) == 1 + for group in y._iter_groups() + ): + task = "classification" + else: + raise ValueError( + "'Recipe.target' must resolve to a single task type" + ) - for task_dispatcher in task_dispatchers: - task_dispatcher._task = task + for task_dispatcher in task_dispatchers: + task_dispatcher._task = task self._related_processors = None self._related_locations = None @@ -135,11 +165,6 @@ def fit_transform( ensemble_table, generator=generator, ) - if len(related_ensembles[name]) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) for module in _iter_modules(processor): if isinstance(module, sp.SeriesEnsembleProcessor): module._set_buffer( @@ -160,6 +185,7 @@ def fit_transform( x_context_only: EnsembleTable | None = None if context_only_columns: + assert y is not None if not isinstance(x, TableTensor): x = TableTensor.from_tensor(cast(Tensor, x)) context_only_columns = [str(col) for col in context_only_columns] @@ -170,11 +196,6 @@ def fit_transform( ), generator=generator, ) - if len(x_context_only) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) x = x.drop_columns(context_only_columns) inputs = _to_ensemble_table(x, num_members) @@ -182,10 +203,6 @@ def fit_transform( inputs, generator=generator, ) - if len(x) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of ensemble members" - ) for module in series_processors: module._set_buffer( @@ -197,8 +214,8 @@ def fit_transform( ), ) - members: list[MemberContext] = [] - for member_id in range(self.num_members): + members: list[MemberContext | MemberQuery] = [] + for member_id in range(len(x)): related_tables_i: RelatedTables[TableTensor] | None = None if related_tables is not None: related_tables_i = RelatedTables( @@ -209,8 +226,8 @@ def fit_transform( relationships=related_tables.relationships, task_links=related_tables.task_links, ) - members.append( - MemberContext( + if y is not None: + member = MemberContext( x=x[member_id], y=y[member_id], x_context_only=x_context_only[member_id] @@ -219,7 +236,12 @@ def fit_transform( related_tables=related_tables_i, input_stypes=inputs[member_id].stypes, ) - ) + else: + member = MemberQuery( + x=x[member_id], + related_tables=related_tables_i, + ) + members.append(member) return tuple(members) @@ -231,10 +253,6 @@ def transform( """Transform query data.""" x: EnsembleTable = _to_ensemble_table(x, self._num_estimators) x = self.recipe.features.transform_ensemble(x) - if len(x) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of ensemble members" - ) related_ensembles: Mapping[str, EnsembleTable] = {} if related_tables is not None: @@ -252,14 +270,9 @@ def transform( related_ensembles[name] = processor.transform_ensemble( ensemble_table ) - if len(related_ensembles[name]) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) members: list[MemberQuery] = [] - for member_id in range(self.num_members): + for member_id in range(len(x)): related_tables_i: RelatedTables[TableTensor] | None = None if related_tables is not None: related_tables_i = RelatedTables( From 1feafb1d0cffcae665710ae0d25be241e36ffbb1 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 08:39:35 +0000 Subject: [PATCH 09/15] update --- sdm/models/kumo/timeseries/anomaly/model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 20a094b27..15b445e7c 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -152,14 +152,14 @@ def forward( generator=generator, ) - for callback in callbacks: - out = callback.on_model_forward_end(self, out) - 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"), From 6f18bb41d66fbc1b381319ebd5181e7b49df62cc Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 08:42:04 +0000 Subject: [PATCH 10/15] update --- sdm/processing/numerical/pca.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sdm/processing/numerical/pca.py b/sdm/processing/numerical/pca.py index c1e37e9ee..fee8df351 100644 --- a/sdm/processing/numerical/pca.py +++ b/sdm/processing/numerical/pca.py @@ -63,3 +63,9 @@ def _transform(self, table: TableTensor) -> TableTensor: TableTensor, torch.cat([table.drop_stypes(Stype.numerical), out], dim=-1), ) + + def __repr__(self, *, indent: int = 0) -> str: + return ( + f"{' ' * indent}{self.__class__.__name__}(" + f"num_components={self.num_components})" + ) From b482571257f85ca549c7181c62ec261310a8c347 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 09:02:54 +0000 Subject: [PATCH 11/15] update --- sdm/models/kumo/timeseries/anomaly/model.py | 22 +++++++++------------ 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index f707e208d..4b9d30737 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -124,9 +124,6 @@ def forward( requires_grad = self.training requires_grad |= any(callback.requires_grad for callback in callbacks) - param = next(self.parameters()) - x = x.to(param.device) - recipe_execution = RecipeExecution( self.default_recipe() if recipe is None else copy.deepcopy(recipe) ) @@ -152,16 +149,10 @@ def forward( ) self._validate_feature_stype(query.x) - if query.x.numerical.size(-1) > self.model.num_features: - raise ValueError( - f"{self.__class__.__name__!r} supports at most " - f"{self.model.num_features} features after pre-processing " - f"(got {query.x.numerical.size(-1)})" - ) with inference_mode("grad" if requires_grad else "inference"): out = self.model( - x=query.x.numerical.to(param.dtype), + x=query.x.numerical, num_samples=num_samples, num_steps=num_steps, generator=generator, @@ -181,9 +172,7 @@ def forward( torch.amp.autocast(x.device.type, enabled=False), inference_mode("grad" if requires_grad else "inference"), ): - out = recipe_execution.transform_output([out]) - # Output reducers may already have removed the single-member axis. - return out[0] if out.ndim == query.x.ndim + 1 else out + return recipe_execution.transform_output([out])[0] class _KumoAnomaly(torch.nn.Module): @@ -236,6 +225,13 @@ def forward( generator: torch.Generator | None, ) -> Tensor: # [..., R] + if x.size(-1) > self.num_features: + raise ValueError( + f"{self.__class__.__name__!r} supports at most " + f"{self.num_features} features after pre-processing " + f"(got {x.size(-1)})" + ) + *size, length, features = x.size() x = x.reshape(-1, length, features) x = F.pad(x, (0, self.num_features - features)) From eab7ee099a798c3025bfdc6229e106bef6dfb49e Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 09:06:49 +0000 Subject: [PATCH 12/15] update --- sdm/models/kumo/timeseries/anomaly/model.py | 6 +- sdm/processing/execution.py | 79 +++++++-------------- test/processing/test_execution.py | 25 +------ 3 files changed, 29 insertions(+), 81 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index 4b9d30737..bba0225af 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -178,7 +178,7 @@ def forward( class _KumoAnomaly(torch.nn.Module): def __init__( self, - num_features: int = 40, + in_channels: int = 40, time_channels: int = 256, feature_channels: int = 64, channels: int = 128, @@ -195,14 +195,14 @@ def __init__( super().__init__() factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} - self.num_features = num_features + self.in_channels = in_channels self.num_steps = num_steps self.window_size = window_size self.num_blocks = num_blocks self.batch_size = batch_size self.side_embedding = SideEmbedding( - num_features=num_features, + num_features=in_channels, time_channels=time_channels, feature_channels=feature_channels, **factory_kwargs, diff --git a/sdm/processing/execution.py b/sdm/processing/execution.py index dbc8a5789..68a4b0e08 100644 --- a/sdm/processing/execution.py +++ b/sdm/processing/execution.py @@ -45,14 +45,13 @@ def __init__(self, recipe: Recipe) -> None: Mapping[str, tuple[tuple[int, int], ...]] | None ) = None self._num_estimators: int | None = None - self._num_members: int | None = None self._y_locations: tuple[tuple[int, int], ...] | None = None @property def num_members(self) -> int: r"""The number of fitted members.""" - assert self._num_members is not None - return self._num_members + assert self._y_locations is not None + return len(self._y_locations) @overload def fit_transform( @@ -91,34 +90,17 @@ def fit_transform( """Fit and transform context data.""" self._num_estimators = num_members - inputs = _to_ensemble_table(x, num_members) - self._num_members = len(inputs) - self._y_locations = None - targets = ( - None - if y is None - else _to_ensemble_table(y, num_members, expand=True) - ) - if targets is not None: - self._num_members = len(targets) - - processors = (self.recipe.features,) - if targets is not None: - processors += (self.recipe.target,) - series_processors = tuple( - module - for processor in processors - for module in _iter_modules(processor) - if isinstance(module, sp.SeriesEnsembleProcessor) - ) - for module in series_processors: - module._set_buffer( - name="_offset", - value=torch.zeros( - size=(self.num_members,), - dtype=torch.long, - device=inputs[0].device, - ), + # Transform target first to be able to resolve task type. Inverse + # target transforms require distinct member assignment: + series_processors: tuple[sp.SeriesEnsembleProcessor, ...] = () + if y is not None: + y: EnsembleTable = _to_ensemble_table(y, num_members, expand=True) + + series_processors = tuple( + module + for processor in (self.recipe.features, self.recipe.target) + for module in _iter_modules(processor) + if isinstance(module, sp.SeriesEnsembleProcessor) ) for module in series_processors: module._set_buffer( @@ -130,13 +112,10 @@ def fit_transform( ), ) - # Transform target first to resolve task type and inversion layout. - if targets is not None: - targets = self.recipe.target.fit_transform_ensemble( - targets, generator=generator + y = self.recipe.target.fit_transform_ensemble( + y, generator=generator ) - self._y_locations = targets._locations - self._num_members = len(targets) + self._y_locations = y._locations task_dispatchers = tuple( module @@ -145,22 +124,19 @@ def fit_transform( if isinstance(module, sp.TaskDispatch) ) if len(task_dispatchers) > 0: - if any( - group.size(-1) != 1 for group in targets._iter_groups() - ): + if any(group.size(-1) != 1 for group in y._iter_groups()): raise ValueError( "Expected the transformed target to contain exactly " "one column" ) if all( - group.numerical.size(-1) == 1 - for group in targets._iter_groups() + group.numerical.size(-1) == 1 for group in y._iter_groups() ): task = "regression" elif all( group.categorical.size(-1) == 1 - for group in targets._iter_groups() + for group in y._iter_groups() ): task = "classification" else: @@ -209,7 +185,7 @@ def fit_transform( x_context_only: EnsembleTable | None = None if context_only_columns: - assert targets is not None + assert y is not None if not isinstance(x, TableTensor): x = TableTensor.from_tensor(cast(Tensor, x)) context_only_columns = [str(col) for col in context_only_columns] @@ -220,14 +196,9 @@ def fit_transform( ), generator=generator, ) - if len(x_context_only) != self.num_members: - raise ValueError( - "Expected inputs to map to the same number of " - "ensemble members" - ) - inputs = _to_ensemble_table( - x.drop_columns(context_only_columns), num_members - ) + x = x.drop_columns(context_only_columns) + + inputs = _to_ensemble_table(x, num_members) x: EnsembleTable = self.recipe.features.fit_transform_ensemble( inputs, generator=generator, @@ -255,10 +226,10 @@ def fit_transform( relationships=related_tables.relationships, task_links=related_tables.task_links, ) - if targets is not None: + if y is not None: member = MemberContext( x=x[member_id], - y=targets[member_id], + y=y[member_id], x_context_only=x_context_only[member_id] if x_context_only is not None else None, diff --git a/test/processing/test_execution.py b/test/processing/test_execution.py index be5702c07..580bdf1f9 100644 --- a/test/processing/test_execution.py +++ b/test/processing/test_execution.py @@ -13,30 +13,7 @@ Stype, TableTensor, ) -from sdm.processing.execution import MemberQuery, RecipeExecution - - -def test_targetless_series_fit_and_transform() -> None: - execution = RecipeExecution(Recipe(features=sp.Detrend())) - x = torch.arange(4.0).view(4, 1) - members = execution.fit_transform(x, None, None, num_members=2) - assert execution.num_members == 2 - assert all(isinstance(member, MemberQuery) for member in members) - for member in members: - torch.testing.assert_close(member.x.numerical, torch.zeros_like(x)) - query = torch.arange(4.0, 6.0).view(2, 1) - for member in execution.transform(query, None): - torch.testing.assert_close(member.x.numerical, torch.zeros_like(query)) - - -def test_supervised_member_count_mismatch() -> None: - execution = RecipeExecution(Recipe()) - with pytest.raises(ValueError, match="same number of ensemble members"): - execution.fit_transform( - x=torch.zeros(2, 4, 3), - y=torch.zeros(3, 4, 1), - related_tables=None, - ) +from sdm.processing.execution import RecipeExecution def test_sequence_uses_batched_fit_states_for_shared_query() -> None: From 08b7932e613059b729b4d27323ef799d8caa4869 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 09:33:08 +0000 Subject: [PATCH 13/15] update --- .../kumo/timeseries/anomaly/denoiser.py | 1 + .../kumo/timeseries/anomaly/embedding.py | 8 +- sdm/models/kumo/timeseries/anomaly/model.py | 81 ++++++++-------- .../kumo/timeseries/anomaly/test_model.py | 93 +------------------ 4 files changed, 47 insertions(+), 136 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/denoiser.py b/sdm/models/kumo/timeseries/anomaly/denoiser.py index f9e4931c6..dc69ce792 100644 --- a/sdm/models/kumo/timeseries/anomaly/denoiser.py +++ b/sdm/models/kumo/timeseries/anomaly/denoiser.py @@ -29,6 +29,7 @@ def __init__( ) -> None: super().__init__() factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.num_steps = num_steps self.diffusion_embedding = DiffusionEmbedding( num_steps=num_steps, diff --git a/sdm/models/kumo/timeseries/anomaly/embedding.py b/sdm/models/kumo/timeseries/anomaly/embedding.py index 5c2feae86..f9f154bc9 100644 --- a/sdm/models/kumo/timeseries/anomaly/embedding.py +++ b/sdm/models/kumo/timeseries/anomaly/embedding.py @@ -87,7 +87,7 @@ class SideEmbedding(torch.nn.Module): before conversion to the parameter dtype. Args: - num_features: Number of features in checkpoint order. + in_channels: Number of features in checkpoint order. time_channels: Even sinusoidal time embedding width, at least two. feature_channels: Learned feature embedding width. is_unconditional: Append the conditioning mask as the last channel. @@ -97,7 +97,7 @@ class SideEmbedding(torch.nn.Module): def __init__( self, - num_features: int, + in_channels: int, time_channels: int = 128, feature_channels: int = 16, is_unconditional: bool = False, @@ -109,7 +109,7 @@ def __init__( self.time_channels = time_channels self.is_unconditional = is_unconditional self.embed_layer = torch.nn.Embedding( - num_embeddings=num_features, + num_embeddings=in_channels, embedding_dim=feature_channels, **factory_kwargs, ) @@ -129,7 +129,7 @@ def forward(self, time: Tensor, mask: Tensor) -> Tensor: time: Model time coordinates with shape ``[B, L]``. Coordinates are used as supplied, with no normalization or resampling. mask: Conditioning mask with shape ``[B, K, L]`` in checkpoint - feature order, where ``K`` equals ``num_features``. One marks + feature order, where ``K`` equals ``in_channels``. One marks an observed value used for conditioning; zero marks a target. Both inputs must be on the module's device. diff --git a/sdm/models/kumo/timeseries/anomaly/model.py b/sdm/models/kumo/timeseries/anomaly/model.py index bba0225af..b4d7929ea 100644 --- a/sdm/models/kumo/timeseries/anomaly/model.py +++ b/sdm/models/kumo/timeseries/anomaly/model.py @@ -186,23 +186,15 @@ def __init__( num_heads: int = 8, embedding_channels: int = 256, num_steps: int = 500, - window_size: int = 100, - num_blocks: int = 4, - batch_size: int = 32, device: torch.device | str | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} - self.in_channels = in_channels - self.num_steps = num_steps - self.window_size = window_size - self.num_blocks = num_blocks - self.batch_size = batch_size self.side_embedding = SideEmbedding( - num_features=in_channels, + in_channels=in_channels, time_channels=time_channels, feature_channels=feature_channels, **factory_kwargs, @@ -222,27 +214,27 @@ def forward( x: Tensor, # [..., R, D] num_samples: int, num_steps: int, - generator: torch.Generator | None, + window_size: int = 100, + num_blocks: int = 4, + batch_size: int = 32, + generator: torch.Generator | None = None, ) -> Tensor: # [..., R] - if x.size(-1) > self.num_features: + if x.size(-1) > self.in_channels: raise ValueError( - f"{self.__class__.__name__!r} supports at most " - f"{self.num_features} features after pre-processing " - f"(got {x.size(-1)})" + f"'KumoAnomaly' supports at most {self.in_channels} columns " + f"after pre-processing (got {x.size(-1)})" ) - *size, length, features = x.size() - x = x.reshape(-1, length, features) - x = F.pad(x, (0, self.num_features - features)) + *B, R, D = x.size() + x = F.pad(x.reshape(-1, R, D), (0, self.in_channels - D)) - # Pad the last window by repeating the last row: - padding = -length % self.window_size - x = torch.cat([x, x[:, -1:].expand(-1, padding, -1)], dim=1) - windows = x.reshape(-1, self.window_size, self.num_features) + # Pad last window by repeating last row: + x = torch.cat([x, x[:, -1:].expand(-1, -R % window_size, -1)], dim=-2) + window = x.reshape(-1, window_size, self.in_channels) schedule = DiffusionSchedule( - num_steps=self.num_steps, + num_steps=self.denoiser.num_steps, beta_start=0.0001, beta_end=0.01, schedule="linear", @@ -252,35 +244,44 @@ def forward( out = torch.cat( [ - self.reconstruct(batch, sampler, num_samples, generator) - for batch in windows.split(self.batch_size) - ] + self.reconstruct( + x=_x, + sampler=sampler, + num_samples=num_samples, + num_blocks=num_blocks, + generator=generator, + ) + for _x in window.split(batch_size) + ], + dim=0, ) - out = out.reshape(x.size())[:, :length, :features] - out = (out - x[:, :length, :features]).abs().mean(dim=-1) - return out.reshape(*size, length) + + out = out.view(x.size())[:, :R, :D] + out = (out - x[:, :R, :D]).abs().mean(dim=-1) + return out.view(*B, R) def reconstruct( self, - x: Tensor, # [B, L, K] + x: Tensor, # [B, R, D] sampler: DPMSolver, num_samples: int, + num_blocks: int, generator: torch.Generator | None, - ) -> Tensor: # [B, L, K] + ) -> Tensor: # [B, R, D] - x = x.transpose(1, 2) # [B, K, L] - batch, _, length = x.size() + B, R, _ = x.size() + x = x.transpose(1, 2) # [B, D, R] mean = x.mean(dim=(1, 2), keepdim=True) std = x.std(dim=(1, 2), keepdim=True) + 1e-5 x = ((x - mean) / std).clamp(-10, 10) - time = torch.arange(length, device=x.device).expand(batch, -1) - blocks = time // (length // self.num_blocks) + time = torch.arange(R, device=x.device).expand(B, -1) + blocks = time // (R // num_blocks) out = x.new_zeros((num_samples, *x.size())) for strategy in range(2): mask = ((blocks % 2) == strategy).to(x.dtype) - mask = mask.unsqueeze(1).expand_as(x) # [B, K, L] + mask = mask.unsqueeze(1).expand_as(x) # [B, D, R] predict_noise = partial( self.predict_noise, cond=mask * x, @@ -297,14 +298,14 @@ def reconstruct( def predict_noise( self, - x: Tensor, # [B, K, L] + x: Tensor, # [B, D, R] step: Tensor, # [1] *, - cond: Tensor, # [B, K, L] - target_mask: Tensor, # [B, K, L] - side_info: Tensor, # [B, S, K, L] + cond: Tensor, # [B, D, R] + target_mask: Tensor, # [B, D, R] + side_info: Tensor, # [B, S, D, R] strategy: Tensor, # [1] - ) -> Tensor: # [B, K, L] + ) -> Tensor: # [B, D, R] return self.denoiser( x=torch.stack([cond, target_mask * x], dim=1), side_info=side_info, diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py index 03ef8f421..fbe71cee4 100644 --- a/test/models/kumo/timeseries/anomaly/test_model.py +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -1,13 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import pytest import torch -import sdm.processing as sp -from sdm import CategoricalTensor, Recipe, Stype, TableTensor -from sdm.explain.gradient import CaptureInputs, EnableInputGradients -from sdm.models import Callback, KumoAnomaly +from sdm.models import KumoAnomaly from sdm.testing import withCUDA @@ -28,90 +24,3 @@ def generator() -> torch.Generator: again = model(x, num_samples=1, num_steps=1, generator=generator()) assert torch.equal(again.numerical, out.numerical) - - -def test_wide_float64_forward() -> None: - model = KumoAnomaly(pretrained=False) - x = torch.randn(60, 50, dtype=torch.float64) - - out = model(x, num_samples=1, num_steps=1) - assert out.size() == (60, 1) - assert out.numerical.dtype == torch.float64 - - -@pytest.mark.parametrize("batch_shape", [(), (2,)]) -@pytest.mark.parametrize("reduce_members", [False, True]) -def test_output_recipe( - batch_shape: tuple[int, ...], reduce_members: bool -) -> None: - model = KumoAnomaly(pretrained=False) - x = torch.randn(*batch_shape, 12, 2) - generator = torch.Generator() - state = generator.get_state() - expected = model(x, num_samples=1, num_steps=1, generator=generator) - recipe = model.default_recipe() - if reduce_members: - recipe.output = sp.AverageEstimators() - recipe.append_output( - lambda table: table.replace_blocks(numerical=table.numerical + 2) - ) - generator.set_state(state) - out = model( - x=x, - recipe=recipe, - num_samples=1, - num_steps=1, - generator=generator, - ) - assert out.numerical_columns == ("mae",) - torch.testing.assert_close(out.numerical, expected.numerical + 2) - - -class DoubleScores(Callback): - def on_model_forward_end( - self, - model: torch.nn.Module, - out: TableTensor, - ) -> TableTensor: - return out.replace_blocks(numerical=2 * out.numerical.double()) - - -def test_callbacks() -> None: - model = KumoAnomaly(pretrained=False) - x = torch.randn(30, 2) - capture = CaptureInputs() - callbacks = [EnableInputGradients(), capture, DoubleScores()] - - def generator() -> torch.Generator: - return torch.Generator().manual_seed(0) - - out = model(x, num_samples=1, num_steps=1, generator=generator()) - doubled = model( - x, - num_samples=1, - num_steps=1, - callbacks=callbacks, - generator=generator(), - ) - assert doubled.numerical.dtype == torch.float32 - torch.testing.assert_close(doubled.numerical, 2 * out.numerical) - - doubled.numerical.sum().backward() - ((x, _),) = capture.inputs - assert x.numerical.grad is not None - - -def test_invalid_inputs() -> None: - model = KumoAnomaly(pretrained=False) - x = TableTensor( - columns={Stype.numerical: ["a"], Stype.categorical: ["b"]}, - numerical=torch.randn(10, 1), - categorical=CategoricalTensor.from_tensor( - torch.randint(0, 5, (10, 1)) - ), - ) - - with pytest.warns(RuntimeWarning, match="unsupported feature stypes"): - model(x, num_samples=1, num_steps=1) - with pytest.raises(ValueError, match="at most 40 features"): - model(torch.randn(10, 41), recipe=Recipe()) From 43e060801571bdf137cd690b2bd988e915e3163f Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 09:35:24 +0000 Subject: [PATCH 14/15] update --- .../timeseries/anomaly/test_side_embedding.py | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/test/models/kumo/timeseries/anomaly/test_side_embedding.py b/test/models/kumo/timeseries/anomaly/test_side_embedding.py index 0c12b4cdf..76961164e 100644 --- a/test/models/kumo/timeseries/anomaly/test_side_embedding.py +++ b/test/models/kumo/timeseries/anomaly/test_side_embedding.py @@ -35,9 +35,9 @@ def test_side_embedding( and not torch.cuda.is_bf16_supported() ): pytest.skip("CUDA bfloat16 not supported") - batch, features, length = shape + B, D, R = shape module = SideEmbedding( - num_features=features, + in_channels=D, time_channels=time_channels, feature_channels=4, is_unconditional=is_unconditional, @@ -46,10 +46,10 @@ def test_side_embedding( ) if dtype == torch.float64: module = module.to(dtype=dtype) - time = torch.arange(batch * length, device=device).reshape(batch, length) + time = torch.arange(B * R, device=device).reshape(B, R) time = time.float().square() * 0.25 - 2.5 time[:, -1] = 100000 - mask = torch.arange(batch * features * length, device=device) + mask = torch.arange(B * D * R, device=device) mask = mask.reshape(shape).remainder(2).bool() actual = module(time, mask) @@ -60,7 +60,7 @@ def test_side_embedding( / time_channels, ) phase = time.unsqueeze(-1) * frequencies - time_embed = torch.empty(batch, length, time_channels, device=device) + time_embed = torch.empty(B, R, time_channels, device=device) time_embed[..., 0::2] = phase.sin() time_embed[..., 1::2] = phase.cos() time_embed = F.layer_norm( @@ -71,7 +71,7 @@ def test_side_embedding( eps=module.time_embed_norm.eps, ) feature_embed = F.embedding( - input=torch.arange(features, device=device), + input=torch.arange(D, device=device), weight=module.embed_layer.weight, ) feature_embed = F.layer_norm( @@ -83,19 +83,14 @@ def test_side_embedding( ) expected = torch.cat( [ - time_embed[:, :, None].expand(-1, -1, features, -1), - feature_embed[None, None].expand(batch, length, -1, -1), + time_embed[:, :, None].expand(-1, -1, D, -1), + feature_embed[None, None].expand(B, R, -1, -1), ], dim=-1, ).permute(0, 3, 2, 1) if is_unconditional: expected = torch.cat([expected, mask[:, None].to(dtype)], dim=1) - assert actual.shape == ( - batch, - time_channels + 4 + int(is_unconditional), - features, - length, - ) + assert actual.shape == (B, time_channels + 4 + int(is_unconditional), D, R) assert actual.dtype == dtype assert actual.device == device torch.testing.assert_close(actual, expected) From 4a61d4324e53e9731d2348fa9dc87c025c2a8cc9 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Thu, 22 Oct 2026 09:36:20 +0000 Subject: [PATCH 15/15] update --- test/models/kumo/timeseries/anomaly/test_model.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/test/models/kumo/timeseries/anomaly/test_model.py b/test/models/kumo/timeseries/anomaly/test_model.py index fbe71cee4..849dd2874 100644 --- a/test/models/kumo/timeseries/anomaly/test_model.py +++ b/test/models/kumo/timeseries/anomaly/test_model.py @@ -12,15 +12,9 @@ def test_forward(device: torch.device) -> None: model = KumoAnomaly(pretrained=False, device=device) x = torch.randn(2, 130, 3, device=device) - def generator() -> torch.Generator: - return torch.Generator(device).manual_seed(0) - - out = model(x, num_samples=1, num_steps=1, generator=generator()) + out = model(x, num_samples=1, num_steps=1) assert out.size() == (2, 130, 1) assert out.columns["numerical"] == ("mae",) assert out.device == device assert out.numerical.isfinite().all() assert (out.numerical >= 0).all() - - again = model(x, num_samples=1, num_steps=1, generator=generator()) - assert torch.equal(again.numerical, out.numerical)