Chronos2 vibe

sdm/models/chronos2/__init__.py

@@ -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",
+]

sdm/models/chronos2/block.py

@@ -0,0 +1,74 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# 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, ReLU, Sequential
+
+from sdm.nn import RMSNorm, RotaryEmbedding, TransformerBlock
+
+
+class LayerNorm(RMSNorm):
+    def forward(
+        self, x: Tensor, rope: RotaryEmbedding | None = None
+    ) -> Tensor:
+        # T5 accumulates in float32 and rounds before multiplying half weights.
+        if rope is not None:
+            x = rope(x)
+        variance = x.float().square().mean(-1, keepdim=True)
+        x = x * (variance + (1e-6 if self.eps is None else self.eps)).rsqrt()
+        if self.weight.dtype in {torch.float16, torch.bfloat16}:
+            x = x.to(self.weight.dtype)
+        return self.weight * x
+
+
+class ResidualBlock(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_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, hidden_channels, **factory_kwargs),
+            ReLU(),
+            Linear(hidden_channels, out_channels, **factory_kwargs),
+        )
+        self.res = Linear(in_channels, out_channels, **factory_kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.mlp(x) + self.res(x)
+
+
+class Chronos2TransformerBlock(TransformerBlock):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        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 = LayerNorm(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=None,
+            scale=1.0,
+            bias=False,
+            **factory_kwargs,
+        )

sdm/models/chronos2/ckpt.py

@@ -0,0 +1,53 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# 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.mlps.{layer}"
+        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

sdm/models/chronos2/icl.py

@@ -0,0 +1,90 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# 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, LayerNorm
+from sdm.nn import 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,
+                rope=rope,
+                **factory_kwargs,
+            )
+            for _ in range(num_layers)
+        )
+        self.group_blocks = ModuleList(
+            Chronos2TransformerBlock(
+                channels=channels,
+                num_heads=num_heads,
+                **factory_kwargs,
+            )
+            for _ in range(num_layers)
+        )
+        self.mlps = ModuleList(
+            Sequential(
+                LayerNorm(channels, eps=1e-6, **factory_kwargs),
+                Linear(
+                    channels, hidden_channels, bias=False, **factory_kwargs
+                ),
+                ReLU(),
+                Linear(
+                    hidden_channels, channels, bias=False, **factory_kwargs
+                ),
+            )
+            for _ in range(num_layers)
+        )
+        self.norm = LayerNorm(channels, eps=1e-6, **factory_kwargs)
+
+    def forward(
+        self,
+        x: Tensor,  # [S, N, C]
+        observed: Tensor,  # [S, N]
+        group_ids: Tensor,  # [S]
+    ) -> Tensor:
+        # Finite additive masks preserve upstream's fully masked-row behavior.
+        minimum = torch.finfo(x.dtype).min
+        time_mask = (~observed).to(x.dtype)[:, None, None, :] * minimum
+        group_mask = group_ids[:, None] == group_ids[None, :]
+        group_mask = group_mask[..., None] & observed[None, :, :]
+        group_mask = (~group_mask).permute(2, 0, 1).unsqueeze(1)
+        group_mask = group_mask.to(x.dtype) * minimum
+        for time_block, group_block, mlp in zip(
+            self.time_blocks, self.group_blocks, self.mlps
+        ):
+            x = time_block(query=x, attn_mask=time_mask)
+            x = group_block(
+                query=x.transpose(0, 1),
+                attn_mask=group_mask,
+            ).transpose(0, 1)
+            x = x + mlp(x)
+        return self.norm(x)