Merge branch 'main' into processing/asinh

Author: rusty1s

sdm/models/kumo/timeseries/anomaly/block.py

@@ -9,7 +9,7 @@
 from torch.nn import Conv1d, LayerNorm, Linear, TransformerEncoderLayer
 
 
-class ResidualBlock(torch.nn.Module):
+class DiffusionBlock(torch.nn.Module):
     """Kumo-Anomaly temporal/feature attention with gated residual outputs.
 
     Time and feature axes of length one bypass their attention and subsequent

sdm/models/kumo/timeseries/anomaly/denoiser.py

@@ -10,7 +10,7 @@
 from torch import Tensor
 from torch.nn import Embedding, LayerNorm, Linear, ModuleList
 
-from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock
+from sdm.models.kumo.timeseries.anomaly.block import DiffusionBlock
 from sdm.models.kumo.timeseries.anomaly.embedding import DiffusionEmbedding
 
 
@@ -48,7 +48,7 @@ def __init__(
         )
         self.residual_layers = ModuleList(
             [
-                ResidualBlock(
+                DiffusionBlock(
                     channels=channels,
                     side_channels=side_channels,
                     embedding_channels=embedding_channels,

sdm/models/timesfm3/block.py

@@ -15,40 +15,17 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# ruff: noqa: D101, D102
+# ruff: noqa: D101
 
 import math
 from typing import Any
 
 import torch
-from torch import Tensor
 from torch.nn import Linear, ReLU, Sequential
 
 from sdm.nn import RMSNorm, RotaryEmbedding, SoftplusScale, TransformerBlock
 
 
-class ResidualBlock(torch.nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        device: torch.device | str | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        super().__init__()
-        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
-
-        self.mlp = Sequential(
-            Linear(in_channels, out_channels, bias=False, **factory_kwargs),
-            ReLU(),
-            Linear(out_channels, out_channels, bias=False, **factory_kwargs),
-        )
-        self.res = Linear(in_channels, out_channels, False, **factory_kwargs)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.mlp(x) + self.res(x)
-
-
 class TimesFM3TransformerBlock(TransformerBlock):
     def __init__(
         self,

sdm/models/timesfm3/model.py

@@ -25,20 +25,20 @@
 
 import torch
 from torch import Tensor
+from torch.nn import ReLU
 
 from sdm import Recipe, RelatedTables, Stype, TableTensor, Task
 from sdm.cache import Cache
 from sdm.models._huggingface import download_checkpoint
 from sdm.models.base import ICLModel
-from sdm.models.timesfm3.block import ResidualBlock
 from sdm.models.timesfm3.ckpt import remap_ckpt
 from sdm.models.timesfm3.icl import ICLBlock
 from sdm.models.timesfm3.recipe import default_recipe
 from sdm.models.timesfm3.util import (
     crossfade_patches,
     gather_future_patches,
 )
-from sdm.nn import RevIN
+from sdm.nn import ResidualMLP, RevIN
 from sdm.tensor.table import TableSchema
 
 
@@ -229,9 +229,12 @@ def __init__(
         self.refinement = refinement
 
         self.revin = RevIN(eps=0, **factory_kwargs)
-        self.patch_embedding = ResidualBlock(
+        self.patch_embedding = ResidualMLP(
             in_channels=2 * (input_patch_size + output_patch_size),
+            hidden_channels=channels,
             out_channels=channels,
+            activation=ReLU(),
+            bias=False,
             **factory_kwargs,
         )
         self.icl_block = ICLBlock(

sdm/models/timestral/block.py

@@ -12,35 +12,11 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
-from torch.nn import Linear, Sequential, Sigmoid, SiLU
+from torch.nn import Linear, Sequential, Sigmoid
 
 from sdm.nn import RMSNorm, RotaryEmbedding, SwiGLU, TransformerBlock
 
 
-class ResidualBlock(torch.nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        hidden_channels: int,
-        out_channels: int,
-        bias: bool = True,
-        device: torch.device | str | None = None,
-        dtype: torch.dtype | None = None,
-    ) -> None:
-        super().__init__()
-        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
-
-        self.mlp = Sequential(
-            Linear(in_channels, hidden_channels, bias, **factory_kwargs),
-            SiLU(),
-            Linear(hidden_channels, out_channels, bias, **factory_kwargs),
-        )
-        self.res = Linear(in_channels, out_channels, bias, **factory_kwargs)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.mlp(x) + self.res(x)
-
-
 class MonotoneQuantiles(torch.nn.Module):
     width: Tensor
 

sdm/nn/__init__.py

@@ -5,6 +5,7 @@
 
 from sdm.nn.rope import RotaryEmbedding
 from sdm.nn.sinusoidal import SinusoidalEmbedding
+from sdm.nn.mlp import ResidualMLP
 from sdm.nn.glu import SwiGLU
 from sdm.nn.softplus import SoftplusScale
 from sdm.nn.rms_norm import RMSNorm
@@ -17,6 +18,7 @@
 __all__ = [
     "RotaryEmbedding",
     "SinusoidalEmbedding",
+    "ResidualMLP",
     "SwiGLU",
     "SoftplusScale",
     "RMSNorm",

sdm/nn/glu.py

@@ -67,12 +67,12 @@ def __init__(
             hidden_channels, channels, bias=bias, **factory_kwargs
         )
 
-    def forward(self, tensor: Tensor) -> Tensor:
+    def forward(self, x: Tensor) -> Tensor:
         r"""The forward pass.
 
         Args:
-            tensor: The tensor.
+            x: The tensor.
         """
-        value = self.up_lin(tensor)
-        gate = F.silu(self.gate_lin(tensor))
+        value = self.up_lin(x)
+        gate = F.silu(self.gate_lin(x))
         return self.down_lin(gate * value)

sdm/nn/mlp.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch.nn import Linear, Sequential
+
+
+class ResidualMLP(torch.nn.Module):
+    r"""Two-layer MLP with a learned linear residual projection.
+
+    .. math::
+
+        \operatorname{MLP}_2(x) + \operatorname{Linear}(x)
+
+    Args:
+        in_channels: The number of input channels.
+        hidden_channels: The number of hidden channels in the MLP.
+        out_channels: The number of output channels.
+        activation: The activation in the MLP.
+        bias: Whether to learn additive biases.
+        device: The device.
+        dtype: The dtype.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        activation: torch.nn.Module,
+        bias: bool = True,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        factory_kwargs: dict[str, Any] = {
+            "bias": bias,
+            "device": device,
+            "dtype": dtype,
+        }
+        self.mlp = Sequential(
+            Linear(in_channels, hidden_channels, **factory_kwargs),
+            activation,
+            Linear(hidden_channels, out_channels, **factory_kwargs),
+        )
+        self.res = Linear(in_channels, out_channels, **factory_kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        r"""The forward pass.
+
+        Args:
+            x: The tensor.
+        """
+        return self.mlp(x) + self.res(x)

test/models/kumo/timeseries/anomaly/test_block.py

@@ -9,7 +9,7 @@
 import torch.nn.functional as F
 from torch import Tensor
 
-from sdm.models.kumo.timeseries.anomaly.block import ResidualBlock
+from sdm.models.kumo.timeseries.anomaly.block import DiffusionBlock
 from sdm.testing import withCUDA
 
 
@@ -34,7 +34,7 @@ def reference_attention(
 
 
 def reference_block(
-    block: ResidualBlock,
+    block: DiffusionBlock,
     x: Tensor,
     side_info: Tensor,
     diffusion_embedding: Tensor,
@@ -79,15 +79,15 @@ def reference_block(
         (3, 5, 1, 2, True),
     ],
 )
-def test_residual_block(
+def test_diffusion_block(
     device: torch.device,
     features: int,
     length: int,
     diffusion_batch: int,
     strategy_batch: int,
     training: bool,
 ) -> None:
-    block = ResidualBlock(
+    block = DiffusionBlock(
         channels=8,
         side_channels=6,
         embedding_channels=12,
@@ -110,8 +110,8 @@ def test_residual_block(
 
 
 @withCUDA
-def test_residual_block_grad_and_state_dict(device: torch.device) -> None:
-    block = ResidualBlock(
+def test_diffusion_block_grad_and_state_dict(device: torch.device) -> None:
+    block = DiffusionBlock(
         channels=8,
         side_channels=6,
         embedding_channels=12,
@@ -135,7 +135,7 @@ def test_residual_block_grad_and_state_dict(device: torch.device) -> None:
         assert value.grad is not None
         assert value.grad.abs().sum() > 0
 
-    restored = ResidualBlock(
+    restored = DiffusionBlock(
         channels=8,
         side_channels=6,
         embedding_channels=12,

test/models/timesfm3/test_block.py test/nn/test_mlp.pytest/models/timesfm3/test_block.py renamed to test/nn/test_mlp.py

@@ -3,16 +3,20 @@
 
 
 import torch
+from torch.nn import ReLU
 
-from sdm.models.timesfm3.block import ResidualBlock
+from sdm.nn import ResidualMLP
 from sdm.testing import withCUDA
 
 
 @withCUDA
-def test_residual_block(device: torch.device) -> None:
-    block = ResidualBlock(
+def test_residual_mlp(device: torch.device) -> None:
+    block = ResidualMLP(
         in_channels=3,
+        hidden_channels=2,
         out_channels=2,
+        activation=ReLU(),
+        bias=False,
         device=device,
     )