From b4d8a993ab7ba80758bafe9f1e03a09b0639141d Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Wed, 14 Oct 2026 08:32:26 -0700 Subject: [PATCH 1/2] Add Kumo-Anomaly diffusion schedule Signed-off-by: Aditi Gautam --- .../kumo/timeseries/anomaly/schedule.py | 84 +++++++++++++ .../kumo/timeseries/anomaly/test_schedule.py | 114 ++++++++++++++++++ 2 files changed, 198 insertions(+) create mode 100644 sdm/models/kumo/timeseries/anomaly/schedule.py create mode 100644 test/models/kumo/timeseries/anomaly/test_schedule.py diff --git a/sdm/models/kumo/timeseries/anomaly/schedule.py b/sdm/models/kumo/timeseries/anomaly/schedule.py new file mode 100644 index 0000000000..628349e5d8 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/schedule.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Literal + +import torch +from torch import Tensor + + +class DiffusionSchedule(torch.nn.Module): + """Discrete noise coefficients for Kumo-Anomaly diffusion. + + The one-dimensional buffers ``beta``, ``alpha_hat = 1 - beta``, and + ``alpha = cumprod(alpha_hat)`` have length ``num_steps``. Only ``alpha`` + is clipped to ``[1e-6, 1 - 1e-6]``, matching Kumo-TS. Tables are computed + in float64 on CPU before conversion to the requested device and dtype. + They are derived from configuration and excluded from the state dict. + + Args: + num_steps: Number of diffusion steps, at least one. + beta_start: Initial noise variance, between zero and one. + beta_end: Final noise variance, between ``beta_start`` and one. + schedule: ``"quad"`` interpolates square roots of the noise variances + before squaring; ``"linear"`` interpolates the variances directly. + device: The device. + dtype: Coefficient dtype. Defaults to float64, as in the upstream + NumPy schedule. Use float32 or float64 to retain the alpha bounds. + """ + + beta: Tensor + alpha_hat: Tensor + alpha: Tensor + + def __init__( + self, + num_steps: int, + beta_start: float, + beta_end: float, + schedule: Literal["quad", "linear"] = "quad", + device: torch.device | str | None = None, + dtype: torch.dtype = torch.float64, + ) -> None: + super().__init__() + self.num_steps = num_steps + if schedule == "quad": + beta_start = beta_start**0.5 + beta_end = beta_end**0.5 + else: + assert schedule == "linear" + beta = torch.linspace( + start=beta_start, + end=beta_end, + steps=num_steps, + device="cpu", + dtype=torch.float64, + ) + if schedule == "quad": + beta = beta.square() + alpha_hat = 1 - beta + alpha = alpha_hat.cumprod(dim=0).clamp(min=1e-6, max=1 - 1e-6) + for name, value in ( + ("beta", beta), + ("alpha_hat", alpha_hat), + ("alpha", alpha), + ): + self.register_buffer( + name=name, + tensor=value.to(device=device, dtype=dtype), + persistent=False, + ) + + def forward(self, step: Tensor) -> Tensor: + """Look up clipped cumulative signal coefficients. + + Args: + step: Integer diffusion indices in ``[0, num_steps)`` with any + shape, on the module's device. + + Returns: + Cumulative ``alpha`` values with the same shape as ``step`` and + the coefficient dtype. Cast to the data dtype and broadcast over + feature/time dimensions when constructing noisy observations. + """ + return self.alpha[step] diff --git a/test/models/kumo/timeseries/anomaly/test_schedule.py b/test/models/kumo/timeseries/anomaly/test_schedule.py new file mode 100644 index 0000000000..55d87d6c9e --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_schedule.py @@ -0,0 +1,114 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Literal + +import numpy as np +import pytest +import torch + +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize( + ("schedule", "num_steps", "beta_start", "beta_end", "dtype"), + [ + ("quad", 1000, 1e-4, 0.5, torch.float64), + ("linear", 1000, 1e-4, 0.02, torch.float64), + ("quad", 1000, 1e-4, 0.5, torch.float32), + ("linear", 1000, 1e-4, 0.02, torch.float32), + ("quad", 1, 1e-4, 0.5, torch.float64), + ("linear", 1, 1e-4, 0.5, torch.float64), + ("quad", 100, 1e-10, 0.9, torch.float64), + ("linear", 100, 1e-10, 0.9, torch.float64), + ], +) +def test_diffusion_schedule( + device: torch.device, + schedule: Literal["quad", "linear"], + num_steps: int, + beta_start: float, + beta_end: float, + dtype: torch.dtype, +) -> None: + """Match the upstream NumPy coefficients and cumulative-only clipping.""" + module = DiffusionSchedule( + num_steps=num_steps, + beta_start=beta_start, + beta_end=beta_end, + schedule=schedule, + device=device, + dtype=dtype, + ) + if schedule == "quad": + beta = np.linspace(beta_start**0.5, beta_end**0.5, num_steps) ** 2 + else: + beta = np.linspace(beta_start, beta_end, num_steps) + alpha_hat = 1 - beta + alpha = np.clip(np.cumprod(alpha_hat), 1e-6, 1 - 1e-6) + for actual, values in ( + (module.beta, beta), + (module.alpha_hat, alpha_hat), + (module.alpha, alpha), + ): + expected = torch.tensor(values, device=device, dtype=dtype) + assert actual.shape == (num_steps,) + assert actual.dtype == dtype + assert actual.device == device + assert not actual.requires_grad + torch.testing.assert_close( + actual=actual, + expected=expected, + rtol=1e-12 if dtype == torch.float64 else 1e-6, + atol=1e-15, + ) + + # Unsorted and repeated indices retain their layout, including endpoints. + step = torch.tensor( + [[num_steps - 1, 0], [num_steps // 2, 0]], device=device + ) + expected = torch.tensor(alpha, device=device, dtype=dtype)[step] + torch.testing.assert_close(module(step), expected) + torch.testing.assert_close(module(step[0, 0]), expected[0, 0]) + empty_step = torch.empty(0, dtype=torch.long, device=device) + assert module(empty_step).shape == (0,) + if beta_start < 1e-6: + assert module.alpha[0] == 1 - 1e-6 + assert module.alpha[-1] == 1e-6 + assert module.alpha_hat[0] > 1 - 1e-6 + + +@withCUDA +def test_diffusion_schedule_conversion(device: torch.device) -> None: + """Convert all derived coefficients without adding checkpoint entries.""" + module = DiffusionSchedule(num_steps=10, beta_start=1e-4, beta_end=0.5) + assert module.beta.dtype == torch.float64 + original = [module.beta, module.alpha_hat, module.alpha] + module = module.to(device=device, dtype=torch.float32) + for actual, expected in zip( + (module.beta, module.alpha_hat, module.alpha), original, strict=True + ): + torch.testing.assert_close(actual, expected.to(device).float()) + assert not module.state_dict() + + +@withCUDA +def test_diffusion_schedule_forward_noise(device: torch.device) -> None: + """Reproduce the float32 coefficients used by upstream forward noising.""" + module = DiffusionSchedule( + num_steps=1000, beta_start=1e-4, beta_end=0.5, device=device + ) + step = torch.tensor([0, 399, 999], device=device) + data = torch.arange(18, device=device).reshape(3, 2, 3).float().sin() + noise = data.cos() * 0.999 + alpha = module(step).float()[:, None, None] + actual = (alpha + 1e-8).sqrt() * data + (1 - alpha + 1e-8).sqrt() * noise + + beta = np.linspace(1e-4**0.5, 0.5**0.5, 1000) ** 2 + alpha = np.clip(np.cumprod(1 - beta), 1e-6, 1 - 1e-6) + alpha = torch.tensor(alpha, device=device, dtype=torch.float32) + alpha = alpha[step][:, None, None] + expected = (alpha + 1e-8).sqrt() * data + (1 - alpha + 1e-8).sqrt() * noise + torch.testing.assert_close(actual, expected, rtol=0, atol=0) From 548c8532c05289a7820735052b5ecf80f1c3ddc2 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:37:35 +0900 Subject: [PATCH 2/2] Honor the default device in Kumo diffusion schedules Keep tests focused on coefficient parity, lookup and buffer conversion. --- .../kumo/timeseries/anomaly/schedule.py | 1 + .../kumo/timeseries/anomaly/test_schedule.py | 34 +++++-------------- 2 files changed, 9 insertions(+), 26 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/schedule.py b/sdm/models/kumo/timeseries/anomaly/schedule.py index 628349e5d8..c7e1d983ac 100644 --- a/sdm/models/kumo/timeseries/anomaly/schedule.py +++ b/sdm/models/kumo/timeseries/anomaly/schedule.py @@ -41,6 +41,7 @@ def __init__( dtype: torch.dtype = torch.float64, ) -> None: super().__init__() + device = torch.get_default_device() if device is None else device self.num_steps = num_steps if schedule == "quad": beta_start = beta_start**0.5 diff --git a/test/models/kumo/timeseries/anomaly/test_schedule.py b/test/models/kumo/timeseries/anomaly/test_schedule.py index 55d87d6c9e..21e67ee870 100644 --- a/test/models/kumo/timeseries/anomaly/test_schedule.py +++ b/test/models/kumo/timeseries/anomaly/test_schedule.py @@ -18,10 +18,7 @@ ("quad", 1000, 1e-4, 0.5, torch.float64), ("linear", 1000, 1e-4, 0.02, torch.float64), ("quad", 1000, 1e-4, 0.5, torch.float32), - ("linear", 1000, 1e-4, 0.02, torch.float32), ("quad", 1, 1e-4, 0.5, torch.float64), - ("linear", 1, 1e-4, 0.5, torch.float64), - ("quad", 100, 1e-10, 0.9, torch.float64), ("linear", 100, 1e-10, 0.9, torch.float64), ], ) @@ -33,7 +30,6 @@ def test_diffusion_schedule( beta_end: float, dtype: torch.dtype, ) -> None: - """Match the upstream NumPy coefficients and cumulative-only clipping.""" module = DiffusionSchedule( num_steps=num_steps, beta_start=beta_start, @@ -71,9 +67,6 @@ def test_diffusion_schedule( ) expected = torch.tensor(alpha, device=device, dtype=dtype)[step] torch.testing.assert_close(module(step), expected) - torch.testing.assert_close(module(step[0, 0]), expected[0, 0]) - empty_step = torch.empty(0, dtype=torch.long, device=device) - assert module(empty_step).shape == (0,) if beta_start < 1e-6: assert module.alpha[0] == 1 - 1e-6 assert module.alpha[-1] == 1e-6 @@ -82,7 +75,6 @@ def test_diffusion_schedule( @withCUDA def test_diffusion_schedule_conversion(device: torch.device) -> None: - """Convert all derived coefficients without adding checkpoint entries.""" module = DiffusionSchedule(num_steps=10, beta_start=1e-4, beta_end=0.5) assert module.beta.dtype == torch.float64 original = [module.beta, module.alpha_hat, module.alpha] @@ -94,21 +86,11 @@ def test_diffusion_schedule_conversion(device: torch.device) -> None: assert not module.state_dict() -@withCUDA -def test_diffusion_schedule_forward_noise(device: torch.device) -> None: - """Reproduce the float32 coefficients used by upstream forward noising.""" - module = DiffusionSchedule( - num_steps=1000, beta_start=1e-4, beta_end=0.5, device=device - ) - step = torch.tensor([0, 399, 999], device=device) - data = torch.arange(18, device=device).reshape(3, 2, 3).float().sin() - noise = data.cos() * 0.999 - alpha = module(step).float()[:, None, None] - actual = (alpha + 1e-8).sqrt() * data + (1 - alpha + 1e-8).sqrt() * noise - - beta = np.linspace(1e-4**0.5, 0.5**0.5, 1000) ** 2 - alpha = np.clip(np.cumprod(1 - beta), 1e-6, 1 - 1e-6) - alpha = torch.tensor(alpha, device=device, dtype=torch.float32) - alpha = alpha[step][:, None, None] - expected = (alpha + 1e-8).sqrt() * data + (1 - alpha + 1e-8).sqrt() * noise - torch.testing.assert_close(actual, expected, rtol=0, atol=0) +def test_diffusion_schedule_default_device() -> None: + with torch.device("meta"): + module = DiffusionSchedule(num_steps=3, beta_start=0.1, beta_end=0.3) + cpu_module = DiffusionSchedule( + num_steps=3, beta_start=0.1, beta_end=0.3, device="cpu" + ) + assert all(buffer.device.type == "meta" for buffer in module.buffers()) + assert all(buffer.device.type == "cpu" for buffer in cpu_module.buffers())