From 915d45b64ba6f14e8a22b8193a7aa046462d5d07 Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Tue, 13 Oct 2026 09:40:32 -0700 Subject: [PATCH 1/2] Add Kumo-Anomaly side-information embedding Signed-off-by: Aditi Gautam --- .../kumo/timeseries/anomaly/embedding.py | 87 ++++++++++++ .../timeseries/anomaly/test_side_embedding.py | 131 ++++++++++++++++++ 2 files changed, 218 insertions(+) create mode 100644 test/models/kumo/timeseries/anomaly/test_side_embedding.py diff --git a/sdm/models/kumo/timeseries/anomaly/embedding.py b/sdm/models/kumo/timeseries/anomaly/embedding.py index f3fca3b6a5..d0e30e50ae 100644 --- a/sdm/models/kumo/timeseries/anomaly/embedding.py +++ b/sdm/models/kumo/timeseries/anomaly/embedding.py @@ -76,3 +76,90 @@ def forward(self, step: Tensor) -> Tensor: """ x = F.silu(self.projection1(self.embedding[step])) return F.silu(self.projection2(x)) + + +class SideEmbedding(torch.nn.Module): + """Embed time coordinates and feature identities for Kumo-Anomaly. + + Sinusoidal time embeddings and learned feature embeddings are normalized + separately before concatenation. The unconditional model also appends the + conditioning mask, matching Kumo-TS. Time phases are computed in float32 + before conversion to the parameter dtype. + + Args: + num_features: 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. + device: The device. + dtype: The parameter and output dtype. + """ + + def __init__( + self, + num_features: int, + time_channels: int = 128, + feature_channels: int = 16, + is_unconditional: bool = False, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + super().__init__() + factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype} + self.time_channels = time_channels + self.is_unconditional = is_unconditional + self.embed_layer = torch.nn.Embedding( + num_embeddings=num_features, + embedding_dim=feature_channels, + **factory_kwargs, + ) + self.embed_norm = torch.nn.LayerNorm( + feature_channels, **factory_kwargs + ) + self.time_embed_norm = torch.nn.LayerNorm( + time_channels, **factory_kwargs + ) + + def forward(self, time: Tensor, mask: Tensor) -> Tensor: + """Build side information without changing time or feature order. + + Args: + 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 + an observed value used for conditioning; zero marks a target. + Both inputs must be on the module's device. + + Returns: + Side information with shape ``[B, S, K, L]`` and parameter dtype. + ``S`` is ``time_channels + feature_channels``, plus one when + ``is_unconditional`` is set. + """ + batch, features, length = mask.shape + frequencies = 1 / torch.pow( + 10000.0, + torch.arange( + 0, + self.time_channels, + 2, + device=time.device, + dtype=torch.float32, + ) + / self.time_channels, + ) + angles = time.float().unsqueeze(-1) * frequencies + time_embed = torch.stack([angles.sin(), angles.cos()], dim=-1) + time_embed = self.time_embed_norm( + time_embed.flatten(-2).to(self.embed_layer.weight.dtype) + ) + time_embed = time_embed.unsqueeze(2).expand(-1, -1, features, -1) + feature_embed = self.embed_norm(self.embed_layer.weight) + feature_embed = feature_embed[None, None].expand(batch, length, -1, -1) + side_info = torch.cat([time_embed, feature_embed], dim=-1) + side_info = side_info.permute(0, 3, 2, 1) # [B, S, K, L] + if self.is_unconditional: + side_info = torch.cat( + [side_info, mask.unsqueeze(1).to(side_info.dtype)], dim=1 + ) + return side_info diff --git a/test/models/kumo/timeseries/anomaly/test_side_embedding.py b/test/models/kumo/timeseries/anomaly/test_side_embedding.py new file mode 100644 index 0000000000..5176a6a13b --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_side_embedding.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.nn.functional as F + +from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize("is_unconditional", [False, True]) +@pytest.mark.parametrize("time_channels", [8, 128]) +@pytest.mark.parametrize("shape", [(1, 1, 1), (2, 3, 7)]) +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float64, torch.float16, torch.bfloat16] +) +def test_side_embedding( + device: torch.device, + is_unconditional: bool, + time_channels: int, + shape: tuple[int, int, int], + dtype: torch.dtype, +) -> None: + """Match Kumo-TS time/feature conditioning, mask order, and gradients.""" + if ( + device.type == "cuda" + and dtype == torch.bfloat16 + and not torch.cuda.is_bf16_supported() + ): + pytest.skip("CUDA bfloat16 not supported") + batch, features, length = shape + module = SideEmbedding( + num_features=features, + time_channels=time_channels, + feature_channels=4, + is_unconditional=is_unconditional, + device=device, + dtype=dtype, + ) + time = torch.arange(batch * length, device=device).reshape(batch, length) + time = time.float().square() * 0.25 - 2.5 + mask = torch.arange(batch * features * length, device=device) + mask = mask.reshape(shape).remainder(2).bool() + actual = module(time, mask) + + # Upstream interleaves sin/cos, then normalizes time and feature channels. + frequencies = 1 / torch.pow( + 10000.0, + torch.arange(0, time_channels, 2, device=device) / time_channels, + ) + phase = time.unsqueeze(-1) * frequencies + time_embed = torch.empty(batch, length, time_channels, device=device) + time_embed[..., 0::2] = phase.sin() + time_embed[..., 1::2] = phase.cos() + time_embed = F.layer_norm( + input=time_embed.to(dtype), + normalized_shape=(time_channels,), + weight=module.time_embed_norm.weight, + bias=module.time_embed_norm.bias, + eps=module.time_embed_norm.eps, + ) + feature_embed = F.embedding( + torch.arange(features, device=device), module.embed_layer.weight + ) + feature_embed = F.layer_norm( + input=feature_embed, + normalized_shape=(4,), + weight=module.embed_norm.weight, + bias=module.embed_norm.bias, + eps=module.embed_norm.eps, + ) + expected = torch.cat( + [ + time_embed[:, :, None].expand(-1, -1, features, -1), + feature_embed[None, None].expand(batch, length, -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.dtype == dtype + assert actual.device == device + torch.testing.assert_close(actual, expected) + actual.float().square().sum().backward() + for parameter in module.parameters(): + assert parameter.grad is not None + assert parameter.grad.isfinite().all() + + changed = module(time, ~mask) + if is_unconditional: + torch.testing.assert_close(changed[:, :-1], actual[:, :-1]) + torch.testing.assert_close(changed[:, -1], (~mask).to(dtype)) + else: + torch.testing.assert_close(changed, actual) + + +@withCUDA +def test_side_embedding_conversion_and_state_dict( + device: torch.device, +) -> None: + """Preserve conditioning after dtype conversion and state loading.""" + module = SideEmbedding( + num_features=3, + time_channels=8, + feature_channels=4, + is_unconditional=True, + ).to(device=device, dtype=torch.float64) + time = torch.tensor([[0, 1, 100000]], device=device) + mask = torch.ones(1, 3, 3, device=device, dtype=torch.float32) + expected = module(time, mask) + restored = SideEmbedding( + num_features=3, + time_channels=8, + feature_channels=4, + is_unconditional=True, + device=device, + dtype=torch.float64, + ) + restored.load_state_dict(module.state_dict(), strict=True) + actual = restored(time, mask) + assert actual.dtype == torch.float64 + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(actual[:, -1], mask.double()) From dff361491ba88dd08a6a02f8ae6a1f421cc3cc92 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:07:39 +0900 Subject: [PATCH 2/2] Focus side embedding tests and use keyword arguments Signed-off-by: Ard <88337265+aw471@users.noreply.github.com> --- .../kumo/timeseries/anomaly/embedding.py | 19 +++--- .../timeseries/anomaly/test_side_embedding.py | 66 +++++++------------ 2 files changed, 35 insertions(+), 50 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/embedding.py b/sdm/models/kumo/timeseries/anomaly/embedding.py index d0e30e50ae..5c2feae86a 100644 --- a/sdm/models/kumo/timeseries/anomaly/embedding.py +++ b/sdm/models/kumo/timeseries/anomaly/embedding.py @@ -114,10 +114,12 @@ def __init__( **factory_kwargs, ) self.embed_norm = torch.nn.LayerNorm( - feature_channels, **factory_kwargs + normalized_shape=feature_channels, + **factory_kwargs, ) self.time_embed_norm = torch.nn.LayerNorm( - time_channels, **factory_kwargs + normalized_shape=time_channels, + **factory_kwargs, ) def forward(self, time: Tensor, mask: Tensor) -> Tensor: @@ -138,11 +140,11 @@ def forward(self, time: Tensor, mask: Tensor) -> Tensor: """ batch, features, length = mask.shape frequencies = 1 / torch.pow( - 10000.0, - torch.arange( - 0, - self.time_channels, - 2, + self=10000.0, + exponent=torch.arange( + start=0, + end=self.time_channels, + step=2, device=time.device, dtype=torch.float32, ) @@ -160,6 +162,7 @@ def forward(self, time: Tensor, mask: Tensor) -> Tensor: side_info = side_info.permute(0, 3, 2, 1) # [B, S, K, L] if self.is_unconditional: side_info = torch.cat( - [side_info, mask.unsqueeze(1).to(side_info.dtype)], dim=1 + tensors=[side_info, mask.unsqueeze(1).to(side_info.dtype)], + dim=1, ) return side_info diff --git a/test/models/kumo/timeseries/anomaly/test_side_embedding.py b/test/models/kumo/timeseries/anomaly/test_side_embedding.py index 5176a6a13b..0c12b4cdfd 100644 --- a/test/models/kumo/timeseries/anomaly/test_side_embedding.py +++ b/test/models/kumo/timeseries/anomaly/test_side_embedding.py @@ -10,11 +10,17 @@ @withCUDA -@pytest.mark.parametrize("is_unconditional", [False, True]) -@pytest.mark.parametrize("time_channels", [8, 128]) -@pytest.mark.parametrize("shape", [(1, 1, 1), (2, 3, 7)]) @pytest.mark.parametrize( - "dtype", [torch.float32, torch.float64, torch.float16, torch.bfloat16] + ("is_unconditional", "time_channels", "shape", "dtype"), + [ + (False, 8, (2, 3, 7), torch.float32), + (True, 128, (2, 3, 7), torch.float32), + (False, 128, (1, 1, 1), torch.float32), + (True, 8, (1, 1, 1), torch.float32), + (False, 8, (2, 3, 7), torch.float16), + (True, 8, (2, 3, 7), torch.bfloat16), + (True, 8, (2, 3, 7), torch.float64), + ], ) def test_side_embedding( device: torch.device, @@ -23,7 +29,6 @@ def test_side_embedding( shape: tuple[int, int, int], dtype: torch.dtype, ) -> None: - """Match Kumo-TS time/feature conditioning, mask order, and gradients.""" if ( device.type == "cuda" and dtype == torch.bfloat16 @@ -37,18 +42,22 @@ def test_side_embedding( feature_channels=4, is_unconditional=is_unconditional, device=device, - dtype=dtype, + dtype=None if dtype == torch.float64 else dtype, ) + if dtype == torch.float64: + module = module.to(dtype=dtype) time = torch.arange(batch * length, device=device).reshape(batch, length) time = time.float().square() * 0.25 - 2.5 + time[:, -1] = 100000 mask = torch.arange(batch * features * length, device=device) mask = mask.reshape(shape).remainder(2).bool() actual = module(time, mask) # Upstream interleaves sin/cos, then normalizes time and feature channels. frequencies = 1 / torch.pow( - 10000.0, - torch.arange(0, time_channels, 2, device=device) / time_channels, + self=10000.0, + exponent=torch.arange(0, time_channels, 2, device=device) + / time_channels, ) phase = time.unsqueeze(-1) * frequencies time_embed = torch.empty(batch, length, time_channels, device=device) @@ -62,7 +71,8 @@ def test_side_embedding( eps=module.time_embed_norm.eps, ) feature_embed = F.embedding( - torch.arange(features, device=device), module.embed_layer.weight + input=torch.arange(features, device=device), + weight=module.embed_layer.weight, ) feature_embed = F.layer_norm( input=feature_embed, @@ -89,10 +99,11 @@ def test_side_embedding( assert actual.dtype == dtype assert actual.device == device torch.testing.assert_close(actual, expected) - actual.float().square().sum().backward() - for parameter in module.parameters(): - assert parameter.grad is not None - assert parameter.grad.isfinite().all() + if dtype == torch.float64: + actual.square().sum().backward() + for parameter in module.parameters(): + assert parameter.grad is not None + assert parameter.grad.isfinite().all() changed = module(time, ~mask) if is_unconditional: @@ -100,32 +111,3 @@ def test_side_embedding( torch.testing.assert_close(changed[:, -1], (~mask).to(dtype)) else: torch.testing.assert_close(changed, actual) - - -@withCUDA -def test_side_embedding_conversion_and_state_dict( - device: torch.device, -) -> None: - """Preserve conditioning after dtype conversion and state loading.""" - module = SideEmbedding( - num_features=3, - time_channels=8, - feature_channels=4, - is_unconditional=True, - ).to(device=device, dtype=torch.float64) - time = torch.tensor([[0, 1, 100000]], device=device) - mask = torch.ones(1, 3, 3, device=device, dtype=torch.float32) - expected = module(time, mask) - restored = SideEmbedding( - num_features=3, - time_channels=8, - feature_channels=4, - is_unconditional=True, - device=device, - dtype=torch.float64, - ) - restored.load_state_dict(module.state_dict(), strict=True) - actual = restored(time, mask) - assert actual.dtype == torch.float64 - torch.testing.assert_close(actual, expected) - torch.testing.assert_close(actual[:, -1], mask.double())