Enable cross-learning in Chronos2

sdm/models/chronos2/model.py

@@ -50,6 +50,10 @@ class Chronos2(ICLModel):
     Args:
         pretrained: Whether to load the pretrained checkpoint.
         device: The device.
+
+    Additional Forward Parameters:
+        cross_learning (bool): Whether examples in a batch will share
+            information through variate attention. (default: ``False``)
     """
 
     supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset(
@@ -81,13 +85,17 @@ def forward(
         x_context: Tensor | TableTensor,
         y_context: Tensor | TableTensor,
         *args: Any,
+        cross_learning: bool = False,
         **kwargs: Any,
     ) -> TableTensor:
         r""":meta private:"""  # noqa: D415
+        if cross_learning:  # Keep estimator groups separate:
+            kwargs["estimator_batch_size"] = 1
         return super().forward(
             x_context[..., -self.model.context_length :, :],
             y_context[..., -self.model.context_length :, :],
             *args,
+            cross_learning=cross_learning,
             **kwargs,
         )
 
@@ -96,13 +104,17 @@ def fit(
         x: Tensor | TableTensor,
         y: Tensor | TableTensor,
         *args: Any,
+        cross_learning: bool = False,
         **kwargs: Any,
     ) -> None:
         r""":meta private:"""  # noqa: D415
+        if cross_learning:  # Keep estimator groups separate:
+            kwargs["estimator_batch_size"] = 1
         super().fit(
             x[..., -self.model.context_length :, :],
             y[..., -self.model.context_length :, :],
             *args,
+            cross_learning=cross_learning,
             **kwargs,
         )
 
@@ -137,6 +149,7 @@ def _forward(
         related_query_tables: RelatedTables[TableTensor] | None,
         cache: Cache | None,
         generator: torch.Generator | None,
+        cross_learning: bool = False,
         **kwargs: Any,
     ) -> TableTensor:  # [..., R_query, Y * 21]
 
@@ -173,7 +186,13 @@ def _forward(
             query[..., y.size(-1) + x_context_only.size(-1) :, :] = (
                 x_query.numerical.transpose(-2, -1)
             )
+            shape = context.size()[:-1]
+            if cross_learning:
+                context = context.flatten(0, -2)
+                query = query.flatten(0, -2)
             out = self.model(context, query)
+            if cross_learning:
+                out = out.unflatten(-3, shape)
             out = out[..., : y.size(-1), :, :].movedim(-1, -3).flatten(-2)
 
         return TableTensor(

test/models/chronos2/test_model.py

@@ -1,33 +1,41 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import pytest
 import torch
 
 from sdm import TableTensor
 from sdm.models import Chronos2
 
 
-def test_forward() -> None:
+@pytest.mark.parametrize("cross_learning", [False, True])
+def test_forward(cross_learning: bool) -> None:
     model = Chronos2(pretrained=False)
 
     # Past and past-and-future covariates:
     x_context = TableTensor.from_tensor(
-        torch.randn(5, 4),
+        torch.randn(2, 5, 4),
         columns=["x1", "x2", "x3", "x4"],
     )
     x_query = TableTensor.from_tensor(
-        torch.randn(3, 2),
+        torch.randn(2, 3, 2),
         columns=["x2", "x4"],
     )
 
     # Target variates:
     y_context = TableTensor.from_tensor(
-        torch.randn(5, 3),
+        torch.randn(2, 5, 3),
         columns=["y1", "y2", "y3"],
     )
 
-    out = model(x_context, y_context, x_query)
-    assert out.size() == (3, 3 * 21)
+    out = model(
+        x_context=x_context,
+        y_context=y_context,
+        x_query=x_query,
+        num_estimators=1,
+        cross_learning=cross_learning,
+    )
+    assert out.size() == (2, 3, 3 * 21)
     assert out.numerical.isfinite().all()
     assert "y1__q10" in out.numerical_columns
     assert "y2__q50" in out.numerical_columns
@@ -36,6 +44,8 @@ def test_forward() -> None:
     model.fit(
         x=x_context,
         y=y_context,
+        num_estimators=1,
         context_only_columns=["x1", "x3"],
+        cross_learning=cross_learning,
     )
     torch.testing.assert_close(model.predict(x_query).numerical, out.numerical)