Add Kumo-Anomaly side-information embedding

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

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)