From 5afde301b188e3db721224123e6b498255f952bb Mon Sep 17 00:00:00 2001 From: rusty1s Date: Fri, 30 Oct 2026 12:14:50 +0000 Subject: [PATCH 1/2] update --- sdm/models/chronos2/model.py | 23 ++++++++++++++++++ test/models/chronos2/test_model.py | 38 ++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/sdm/models/chronos2/model.py b/sdm/models/chronos2/model.py index 0b4098a96..55bdf5167 100644 --- a/sdm/models/chronos2/model.py +++ b/sdm/models/chronos2/model.py @@ -50,6 +50,13 @@ class Chronos2(ICLModel): Args: pretrained: Whether to load the pretrained checkpoint. device: The device. + + Additional Forward Parameters: + cross_learning (bool): Whether examples in the batch share information + through variate attention. Defaults to ``False``. Predictions + depend on batch composition when enabled. Use ``num_estimators=1`` + for inputs shaped ``[B, R, D]`` to treat ``B`` as examples. + For fitted predictions, pass this option to :meth:`fit`. """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( @@ -81,13 +88,18 @@ 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 when merging batch dimensions. + 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 +108,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: + 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 +153,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 +190,13 @@ def _forward( query[..., y.size(-1) + x_context_only.size(-1) :, :] = ( x_query.numerical.transpose(-2, -1) ) + if cross_learning: + shape = context.size()[:-1] + 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( diff --git a/test/models/chronos2/test_model.py b/test/models/chronos2/test_model.py index 9d05b549a..69d356e7e 100644 --- a/test/models/chronos2/test_model.py +++ b/test/models/chronos2/test_model.py @@ -1,6 +1,7 @@ # 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 @@ -39,3 +40,40 @@ def test_forward() -> None: context_only_columns=["x1", "x3"], ) torch.testing.assert_close(model.predict(x_query).numerical, out.numerical) + + +@pytest.mark.parametrize("estimator_batch_size", [1, None]) +def test_cross_learning(estimator_batch_size: int | None) -> None: + model = Chronos2(pretrained=False) + data = torch.randn(2, 2, 5, 2) # [E, B, R, Y] + context = TableTensor.from_tensor(data) + features = TableTensor(size=(2, 2, 5)) + query = TableTensor(size=(2, 2, 3)) + + out = model( + features, + context, + query, + cross_learning=True, + estimator_batch_size=estimator_batch_size, + ) + + # Group examples into target columns separately for each estimator. + grouped = TableTensor.from_tensor(data.transpose(-3, -2).flatten(-2)) + expected = model( + TableTensor(size=(2, 5)), + grouped, + TableTensor(size=(2, 3)), + estimator_batch_size=estimator_batch_size, + ).numerical + expected = expected.reshape(3, 2, 2 * 21).transpose(0, 1) + assert out.size() == (2, 3, 2 * 21) + torch.testing.assert_close(out.numerical, expected) + + model.fit( + features, + context, + cross_learning=True, + estimator_batch_size=estimator_batch_size, + ) + torch.testing.assert_close(model.predict(query).numerical, out.numerical) From 3986f717083a81d2b874a53a6301b69c028ca6a6 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Fri, 30 Oct 2026 12:20:48 +0000 Subject: [PATCH 2/2] update --- sdm/models/chronos2/model.py | 14 +++----- test/models/chronos2/test_model.py | 58 ++++++++---------------------- 2 files changed, 20 insertions(+), 52 deletions(-) diff --git a/sdm/models/chronos2/model.py b/sdm/models/chronos2/model.py index 55bdf5167..2498c326f 100644 --- a/sdm/models/chronos2/model.py +++ b/sdm/models/chronos2/model.py @@ -52,11 +52,8 @@ class Chronos2(ICLModel): device: The device. Additional Forward Parameters: - cross_learning (bool): Whether examples in the batch share information - through variate attention. Defaults to ``False``. Predictions - depend on batch composition when enabled. Use ``num_estimators=1`` - for inputs shaped ``[B, R, D]`` to treat ``B`` as examples. - For fitted predictions, pass this option to :meth:`fit`. + cross_learning (bool): Whether examples in a batch will share + information through variate attention. (default: ``False``) """ supported_feature_stypes: ClassVar[frozenset[Stype]] = frozenset( @@ -92,8 +89,7 @@ def forward( **kwargs: Any, ) -> TableTensor: r""":meta private:""" # noqa: D415 - if cross_learning: - # Keep estimator groups separate when merging batch dimensions. + if cross_learning: # Keep estimator groups separate: kwargs["estimator_batch_size"] = 1 return super().forward( x_context[..., -self.model.context_length :, :], @@ -112,7 +108,7 @@ def fit( **kwargs: Any, ) -> None: r""":meta private:""" # noqa: D415 - if cross_learning: + if cross_learning: # Keep estimator groups separate: kwargs["estimator_batch_size"] = 1 super().fit( x[..., -self.model.context_length :, :], @@ -190,8 +186,8 @@ def _forward( query[..., y.size(-1) + x_context_only.size(-1) :, :] = ( x_query.numerical.transpose(-2, -1) ) + shape = context.size()[:-1] if cross_learning: - shape = context.size()[:-1] context = context.flatten(0, -2) query = query.flatten(0, -2) out = self.model(context, query) diff --git a/test/models/chronos2/test_model.py b/test/models/chronos2/test_model.py index 69d356e7e..e5ea5b777 100644 --- a/test/models/chronos2/test_model.py +++ b/test/models/chronos2/test_model.py @@ -8,27 +8,34 @@ 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 @@ -37,43 +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) - - -@pytest.mark.parametrize("estimator_batch_size", [1, None]) -def test_cross_learning(estimator_batch_size: int | None) -> None: - model = Chronos2(pretrained=False) - data = torch.randn(2, 2, 5, 2) # [E, B, R, Y] - context = TableTensor.from_tensor(data) - features = TableTensor(size=(2, 2, 5)) - query = TableTensor(size=(2, 2, 3)) - - out = model( - features, - context, - query, - cross_learning=True, - estimator_batch_size=estimator_batch_size, - ) - - # Group examples into target columns separately for each estimator. - grouped = TableTensor.from_tensor(data.transpose(-3, -2).flatten(-2)) - expected = model( - TableTensor(size=(2, 5)), - grouped, - TableTensor(size=(2, 3)), - estimator_batch_size=estimator_batch_size, - ).numerical - expected = expected.reshape(3, 2, 2 * 21).transpose(0, 1) - assert out.size() == (2, 3, 2 * 21) - torch.testing.assert_close(out.numerical, expected) - - model.fit( - features, - context, - cross_learning=True, - estimator_batch_size=estimator_batch_size, - ) - torch.testing.assert_close(model.predict(query).numerical, out.numerical)