rusty1s · Oct 30, 2026 · Oct 29, 2026 · Oct 29, 2026 · Oct 29, 2026 · Oct 30, 2026
@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Chronos-2."""
+
+from sdm.models.chronos2.model import Chronos2
+
+__all__ = [
+    "Chronos2",
+]
@@ -0,0 +1,88 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: D101
+
+from typing import Any
+
+import torch
+from torch.nn import Linear, ReLU, Sequential
+
+from sdm.nn import RMSNorm, RotaryEmbedding, TransformerBlock
+
+
+class Chronos2TransformerBlock(TransformerBlock):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        mlp: torch.nn.Module | None,
+        rope: RotaryEmbedding | None = None,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        norm = RMSNorm(channels, eps=1e-6, **factory_kwargs)
+        super().__init__(
+            channels=channels,
+            num_query_heads=num_heads,
+            query_norm=norm,
+            key_value_norm=norm,
+            query_transform=rope,
+            key_transform=rope,
+            mlp=mlp,
+            scale=1.0,
+            bias=False,
+            **factory_kwargs,
+        )
+
+    def peak_bytes_per_example(
+        self,
+        element_size: int,
+        query_length: int,
+        key_value_length: int | None = None,
+    ) -> int:
+        r""":meta private:"""  # noqa: D415
+        length = max(query_length, key_value_length or 0)
+        factor = 7 if self.mlp is None else 10
+        return factor * length * element_size * self.attn.q_dim
+
+
+if __name__ == "__main__":
+    from sdm.testing.memory import benchmark_transformer_block_memory_peak
+
+    benchmark_transformer_block_memory_peak(
+        block=lambda channels, num_heads: Chronos2TransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            mlp=None,
+        ),
+        channels_and_heads=[(768, 12)],
+    )
+    benchmark_transformer_block_memory_peak(
+        block=lambda channels, num_heads: Chronos2TransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            mlp=Sequential(
+                RMSNorm(channels, eps=1e-6),
+                Linear(channels, 3072, bias=False),
+                ReLU(),
+                Linear(3072, channels, bias=False),
+            ),
+        ),
+        channels_and_heads=[(768, 12)],
+    )
@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch import Tensor
+
+
+def remap_ckpt(ckpt: dict[str, Tensor]) -> dict[str, Tensor]:  # noqa: D103
+    out = {"reg_token": ckpt["shared.weight"][1:2]}
+    for source, target in (
+        ("input_patch_embedding", "patch_embedding"),
+        ("output_patch_embedding", "head"),
+    ):
+        for upstream, local in (
+            ("hidden_layer", "mlp.0"),
+            ("output_layer", "mlp.2"),
+            ("residual_layer", "res"),
+        ):
+            for suffix in ("weight", "bias"):
+                out[f"{target}.{local}.{suffix}"] = ckpt[
+                    f"{source}.{upstream}.{suffix}"
+                ]
+    for layer in range(12):
+        prefix = f"encoder.block.{layer}.layer"
+        for index, stack in ((0, "time_blocks"), (1, "group_blocks")):
+            source = f"{prefix}.{index}"
+            target = f"icl_block.{stack}.{layer}"
+            norm = ckpt[f"{source}.layer_norm.weight"]
+            out[f"{target}.query_norm.weight"] = norm
+            out[f"{target}.key_value_norm.weight"] = norm
+            out[f"{target}.attn.qkv_lin.weight"] = torch.cat(
+                [
+                    ckpt[f"{source}.self_attention.{key}.weight"]
+                    for key in "qkv"
+                ]
+            )
+            out[f"{target}.attn.out_lin.weight"] = ckpt[
+                f"{source}.self_attention.o.weight"
+            ]
+        target = f"icl_block.group_blocks.{layer}.mlp"
+        out[f"{target}.0.weight"] = ckpt[f"{prefix}.2.layer_norm.weight"]
+        out[f"{target}.1.weight"] = ckpt[f"{prefix}.2.mlp.wi.weight"]
+        out[f"{target}.3.weight"] = ckpt[f"{prefix}.2.mlp.wo.weight"]
+    out["icl_block.norm.weight"] = ckpt["encoder.final_layer_norm.weight"]
+
+    freq = torch.arange(0, 64, 2, device=out["reg_token"].device).float()
+    inv_freq = 1.0 / (10_000 ** (freq / 64))
+    for layer in range(12):
+        for transform in ("query_transform", "key_transform"):
+            key = f"icl_block.time_blocks.{layer}.attn.{transform}.inv_freq"
+            out[key] = inv_freq
+    return out
@@ -0,0 +1,104 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: D101, D102
+
+from typing import Any
+
+import torch
+from torch import Tensor
+from torch.nn import Linear, ModuleList, ReLU, Sequential
+
+from sdm.models.chronos2.block import Chronos2TransformerBlock
+from sdm.nn import RMSNorm, RotaryEmbedding
+
+
+class ICLBlock(torch.nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_layers: int,
+        num_heads: int,
+        hidden_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}
+
+        rope = RotaryEmbedding(
+            channels=channels // num_heads,
+            layout="split_half",
+            theta=10_000,
+            requires_grad=False,
+            device=device,
+        )
+        self.time_blocks = ModuleList(
+            Chronos2TransformerBlock(
+                channels=channels,
+                num_heads=num_heads,
+                mlp=None,
+                rope=rope,
+                **factory_kwargs,
+            )
+            for _ in range(num_layers)
+        )
+        self.group_blocks = ModuleList(
+            Chronos2TransformerBlock(
+                channels=channels,
+                num_heads=num_heads,
+                mlp=Sequential(
+                    RMSNorm(channels, eps=1e-6, **factory_kwargs),
+                    Linear(channels, hidden_channels, False, **factory_kwargs),
+                    ReLU(),
+                    Linear(hidden_channels, channels, False, **factory_kwargs),
+                ),
+                **factory_kwargs,
+            )
+            for _ in range(num_layers)
+        )
+        self.norm = RMSNorm(channels, eps=1e-6, **factory_kwargs)
+
+    def forward(
+        self,
+        x: Tensor,  # [..., C, N, D]
+        patch_mask: Tensor | None = None,  # [..., C, N]
+    ) -> Tensor:  # [..., C, N, D]
+
+        time_attn_mask = var_attn_mask = None
+        if patch_mask is not None:
+            time_attn_mask = ~patch_mask.unsqueeze(-2).unsqueeze(-2)
+            var_attn_mask = time_attn_mask.transpose(-4, -1)
+
+        for time_block, group_block in zip(
+            self.time_blocks, self.group_blocks
+        ):
+            x = time_block(
+                query=x,
+                attn_mask=time_attn_mask,
+                batch_size_limit="auto",
+                out=None if torch.is_grad_enabled() else x,
+            )
+            x = x.transpose(-3, -2)
+            x = group_block(
+                query=x,
+                attn_mask=var_attn_mask,
+                batch_size_limit="auto",
+                out=None if torch.is_grad_enabled() else x,
+            ).transpose(-3, -2)
+
+        return self.norm(x)