diff --git a/sdm/models/chronos2/model.py b/sdm/models/chronos2/model.py index 0b4098a96..2498c326f 100644 --- a/sdm/models/chronos2/model.py +++ b/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( diff --git a/test/models/chronos2/test_model.py b/test/models/chronos2/test_model.py index 9d05b549a..e5ea5b777 100644 --- a/test/models/chronos2/test_model.py +++ b/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)