fix: unpack fcsts from batches for chronos 2 (#253)

Author: AzulGarza

tests/models/conftest.py

@@ -62,7 +62,8 @@ def disable_mps_session(monkeypatch):
     ADIDA(),
     Prophet(),
     Chronos(repo_id="amazon/chronos-bolt-tiny", alias="Chronos-Bolt"),
-    Chronos(repo_id="s3://autogluon/chronos-2", alias="Chronos-2"),
+    Chronos(repo_id="amazon/chronos-2", alias="Chronos-2"),
+    Chronos(repo_id="amazon/chronos-2", alias="Chronos-2", batch_size=2),
     FlowState(repo_id="ibm-research/flowstate"),
     FlowState(
         repo_id="ibm-granite/granite-timeseries-flowstate-r1",

timecopilot/models/foundation/chronos.py

@@ -51,7 +51,7 @@ def __init__(
 
             | Model ID                                                               | Parameters |
             | ---------------------------------------------------------------------- | ---------- |
-            | [`s3://autogluon/chronos-2`](https://arxiv.org/abs/2510.15821)   | 120M         |
+            | [`amazon/chronos-2`](https://huggingface.co/amazon/chronos-2)   | 120M         |
             | [`amazon/chronos-bolt-tiny`](https://huggingface.co/amazon/chronos-bolt-tiny)   | 9M         |
             | [`amazon/chronos-bolt-mini`](https://huggingface.co/amazon/chronos-bolt-mini)   | 21M        |
             | [`amazon/chronos-bolt-small`](https://huggingface.co/amazon/chronos-bolt-small) | 48M        |
@@ -118,8 +118,12 @@ def _predict(
             ]  # list of tuples
             fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False)
             if isinstance(model, Chronos2Pipeline):
-                fcsts_mean = fcsts_mean[0]
-                fcsts_quantiles = fcsts_quantiles[0]
+                fcsts_mean = [f_mean for fcst in fcsts_mean for f_mean in fcst]  # type: ignore
+                fcsts_quantiles = [
+                    f_quantile
+                    for fcst in fcsts_quantiles
+                    for f_quantile in fcst  # type: ignore
+                ]
             fcsts_mean_np = torch.cat(fcsts_mean).numpy()
             fcsts_quantiles_np = torch.cat(fcsts_quantiles).numpy()
         else:
@@ -131,7 +135,7 @@ def _predict(
                 for batch in tqdm(dataset)
             ]
             if isinstance(model, Chronos2Pipeline):
-                fcsts = fcsts[0]
+                fcsts = [f_fcst for fcst in fcsts for f_fcst in fcst]  # type: ignore
             fcsts = torch.cat(fcsts)
             if isinstance(model, ChronosPipeline):
                 # for t5 models, `predict` returns a tensor of shape