feat: bump timecopilot tirex dep

pyproject.toml

@@ -96,6 +96,7 @@ dependencies = [
   "timecopilot-granite-tsfm>=0.1.2",
   "timecopilot-timesfm>=0.2.1",
   "timecopilot-tirex>=0.1.0 ; python_full_version >= '3.11'",
+  "timecopilot-tirex>=0.1.1",
   "timecopilot-toto-2>=0.1.1",
   "timecopilot-toto>=0.1.6",
   "timecopilot-uni2ts>=0.1.2 ; python_full_version < '3.14'",

timecopilot/models/foundation/tirex.py

@@ -15,6 +15,8 @@
 from ..utils.forecaster import Forecaster, QuantileConverter
 from .utils import TimeSeriesDataset
 
+DEFAULT_QUANTILES_TIREX = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+
 
 class TiRex(Forecaster):
     """
@@ -86,31 +88,20 @@ def _forecast(
         quantiles: list[float] | None,
     ) -> tuple[np.ndarray, np.ndarray | None]:
         """handles distinction between quantiles and no quantiles"""
-        if quantiles is not None:
-            fcsts = [
-                model.forecast(
-                    batch,
-                    prediction_length=h,
-                    quantile_levels=quantiles,
-                    output_type="numpy",
-                )
-                for batch in tqdm(dataset)
-            ]  # list of tuples
-            fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False)
-            fcsts_quantiles_np = np.concatenate(fcsts_quantiles)
-            fcsts_mean_np = np.concatenate(fcsts_mean)
-        else:
-            fcsts = [
-                model.forecast(
-                    batch,
-                    prediction_length=h,
-                    output_type="numpy",
-                )
-                for batch in tqdm(dataset)
-            ]
-            fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False)
-            fcsts_mean_np = np.concatenate(fcsts_mean)
-            fcsts_quantiles_np = None
+        fcsts = [
+            model.forecast(
+                batch,
+                prediction_length=h,
+                output_type="numpy",
+            )
+            for batch in tqdm(dataset)
+        ]  # list of tuples
+        fcsts_quantiles, fcsts_mean = zip(*fcsts, strict=False)
+        fcsts_mean_np = np.concatenate(fcsts_mean)
+        fcsts_quantiles_np = (
+            None if quantiles is None else np.concatenate(fcsts_quantiles)
+        )
+
         return fcsts_mean_np, fcsts_quantiles_np
 
     def forecast(
@@ -164,12 +155,23 @@ def forecast(
                     - prediction intervals if `level` is specified.
                     - quantile forecasts if `quantiles` is specified.
 
-                For multi-series data, the output retains the same unique
-                identifiers as the input DataFrame.
+                For multi-series data, the output retains the
+                same unique identifiers as the input DataFrame.
         """
         freq = self._maybe_infer_freq(df, freq)
         qc = QuantileConverter(level=level, quantiles=quantiles)
-        dataset = TimeSeriesDataset.from_df(df, batch_size=self.batch_size)
+        if qc.quantiles is not None and len(qc.quantiles) != len(
+            DEFAULT_QUANTILES_TIREX
+        ):
+            raise ValueError(
+                "TiRex only supports the default quantiles, "
+                "please use the default quantiles or default level, "
+            )
+        dataset = TimeSeriesDataset.from_df(
+            df,
+            batch_size=self.batch_size,
+        )
+
         fcst_df = dataset.make_future_dataframe(h=h, freq=freq)
         with self._get_model() as model:
             fcsts_mean_np, fcsts_quantiles_np = self._forecast(