From d28ca9127c1f6529e5ce0d8d657840aa18a2a69b Mon Sep 17 00:00:00 2001
From: rusty1s <mfey@nvidia.com>
Date: Thu, 29 Oct 2026 17:42:57 +0000
Subject: [PATCH 1/2] update

---
 sdm/models/timestral/block.py | 69 ++++++++++++++++++++++++++++++++++-
 1 file changed, 67 insertions(+), 2 deletions(-)

diff --git a/sdm/models/timestral/block.py b/sdm/models/timestral/block.py
index b34752aae9..3243a5e031 100644
--- a/sdm/models/timestral/block.py
+++ b/sdm/models/timestral/block.py
@@ -4,16 +4,81 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# ruff: noqa: D101
+# ruff: noqa: D101, D102
 
+from collections.abc import Sequence
 from typing import Any
 
 import torch
-from torch.nn import Linear, Sequential, Sigmoid
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Linear, Sequential, Sigmoid, SiLU
 
 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
+
+    def __init__(
+        self,
+        levels: Sequence[float],
+        width_factor: float = 3.0,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        self.width_factor = width_factor
+        self.median_index = list(levels).index(0.5)
+
+        widths = [
+            levels[i + 1] - levels[i]
+            if i < self.median_index
+            else levels[i] - levels[i - 1]
+            if i > self.median_index
+            else 0.0
+            for i in range(len(levels))
+        ]
+        self.register_buffer(
+            "width",
+            torch.tensor(widths, device=device, dtype=dtype),
+            persistent=False,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        m = self.median_index
+        median = x[..., m : m + 1]
+
+        inc = self.width_factor * self.width.to(x) * F.softplus(x)
+        up = inc[..., m + 1 :].cumsum(-1)
+        down = inc[..., :m].flip(-1).cumsum(-1).flip(-1)
+        return torch.cat([median - down, median, median + up], dim=-1)
+
+
 class TimestralTransformerBlock(TransformerBlock):
     def __init__(
         self,

From 74a3921952525b449baae90e07ebf0e79a8c232e Mon Sep 17 00:00:00 2001
From: rusty1s <mfey@nvidia.com>
Date: Thu, 29 Oct 2026 17:43:21 +0000
Subject: [PATCH 2/2] update

---
 sdm/models/timestral/block.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/sdm/models/timestral/block.py b/sdm/models/timestral/block.py
index 3243a5e031..6d250359b9 100644
--- a/sdm/models/timestral/block.py
+++ b/sdm/models/timestral/block.py
@@ -76,6 +76,7 @@ def forward(self, x: Tensor) -> Tensor:
         inc = self.width_factor * self.width.to(x) * F.softplus(x)
         up = inc[..., m + 1 :].cumsum(-1)
         down = inc[..., :m].flip(-1).cumsum(-1).flip(-1)
+
         return torch.cat([median - down, median, median + up], dim=-1)