diff --git a/sdm/models/kumo/timeseries/anomaly/embedding.py b/sdm/models/kumo/timeseries/anomaly/embedding.py index f3fca3b6a5..5c2feae86a 100644 --- a/sdm/models/kumo/timeseries/anomaly/embedding.py +++ b/sdm/models/kumo/timeseries/anomaly/embedding.py @@ -76,3 +76,93 @@ 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( + normalized_shape=feature_channels, + **factory_kwargs, + ) + self.time_embed_norm = torch.nn.LayerNorm( + normalized_shape=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( + self=10000.0, + exponent=torch.arange( + start=0, + end=self.time_channels, + step=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( + 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 new file mode 100644 index 0000000000..0c12b4cdfd --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_side_embedding.py @@ -0,0 +1,113 @@ +# 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", "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, + is_unconditional: bool, + time_channels: int, + shape: tuple[int, int, int], + dtype: torch.dtype, +) -> None: + 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=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( + 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) + 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( + input=torch.arange(features, device=device), + weight=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) + 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: + torch.testing.assert_close(changed[:, :-1], actual[:, :-1]) + torch.testing.assert_close(changed[:, -1], (~mask).to(dtype)) + else: + torch.testing.assert_close(changed, actual)