autogluon.timeseries 1.4.1b20251115__py3-none-any.whl → 1.4.1b20251218__py3-none-any.whl

autogluon/timeseries/models/__init__.py

@@ -1,5 +1,5 @@
 from .autogluon_tabular import DirectTabularModel, PerStepTabularModel, RecursiveTabularModel
-from .chronos import ChronosModel
+from .chronos import Chronos2Model, ChronosModel
 from .gluonts import (
     DeepARModel,
     DLinearModel,
@@ -45,6 +45,7 @@ __all__ = [
     "ETSModel",
     "IMAPAModel",
     "ChronosModel",
+    "Chronos2Model",
     "ModelRegistry",
     "NPTSModel",
     "NaiveModel",

autogluon/timeseries/models/abstract/abstract_timeseries_model.py

@@ -4,7 +4,7 @@ import os
 import re
 import time
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Sequence
 
 import pandas as pd
 from typing_extensions import Self
@@ -75,15 +75,15 @@ class TimeSeriesModelBase(ModelBase, ABC):
 
     def __init__(
         self,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        hyperparameters: Optional[dict[str, Any]] = None,
-        freq: Optional[str] = None,
+        path: str | None = None,
+        name: str | None = None,
+        hyperparameters: dict[str, Any] | None = None,
+        freq: str | None = None,
         prediction_length: int = 1,
-        covariate_metadata: Optional[CovariateMetadata] = None,
+        covariate_metadata: CovariateMetadata | None = None,
         target: str = "target",
         quantile_levels: Sequence[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
-        eval_metric: Union[str, TimeSeriesScorer, None] = None,
+        eval_metric: str | TimeSeriesScorer | None = None,
     ):
         self.name = name or re.sub(r"Model$", "", self.__class__.__name__)
 
@@ -102,7 +102,7 @@ class TimeSeriesModelBase(ModelBase, ABC):
         self.target: str = target
         self.covariate_metadata = covariate_metadata or CovariateMetadata()
 
-        self.freq: Optional[str] = freq
+        self.freq: str | None = freq
         self.prediction_length: int = prediction_length
         self.quantile_levels: list[float] = list(quantile_levels)
 
@@ -117,17 +117,21 @@ class TimeSeriesModelBase(ModelBase, ABC):
         else:
             self.must_drop_median = False
 
-        self._oof_predictions: Optional[list[TimeSeriesDataFrame]] = None
+        self._oof_predictions: list[TimeSeriesDataFrame] | None = None
 
         # user provided hyperparameters and extra arguments that are used during model training
         self._hyperparameters, self._extra_ag_args = self._check_and_split_hyperparameters(hyperparameters)
 
-        self.fit_time: Optional[float] = None  # Time taken to fit in seconds (Training data)
-        self.predict_time: Optional[float] = None  # Time taken to predict in seconds (Validation data)
-        self.predict_1_time: Optional[float] = (
-            None  # Time taken to predict 1 row of data in seconds (with batch size `predict_1_batch_size`)
-        )
-        self.val_score: Optional[float] = None  # Score with eval_metric (Validation data)
+        # Time taken to fit in seconds (Training data)
+        self.fit_time: float | None = None
+        # Time taken to predict in seconds, for a single prediction horizon on validation data
+        self.predict_time: float | None = None
+        # Time taken to predict 1 row of data in seconds (with batch size `predict_1_batch_size`)
+        self.predict_1_time: float | None = None
+        # Useful for ensembles, additional prediction time excluding base models. None for base models.
+        self.predict_time_marginal: float | None = None
+        # Score with eval_metric on validation data
+        self.val_score: float | None = None
 
     def __repr__(self) -> str:
         return self.name
@@ -143,9 +147,14 @@ class TimeSeriesModelBase(ModelBase, ABC):
         self.path = path_context
         self.path_root = self.path.rsplit(self.name, 1)[0]
 
+    def cache_oof_predictions(self, predictions: TimeSeriesDataFrame | list[TimeSeriesDataFrame]) -> None:
+        if isinstance(predictions, TimeSeriesDataFrame):
+            predictions = [predictions]
+        self._oof_predictions = predictions
+
     @classmethod
     def _check_and_split_hyperparameters(
-        cls, hyperparameters: Optional[dict[str, Any]] = None
+        cls, hyperparameters: dict[str, Any] | None = None
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """Given the user-specified hyperparameters, split into `hyperparameters` and `extra_ag_args`, intended
         to be used during model initialization.
@@ -180,7 +189,7 @@ class TimeSeriesModelBase(ModelBase, ABC):
             )
         return hyperparameters, extra_ag_args
 
-    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
+    def save(self, path: str | None = None, verbose: bool = True) -> str:
         if path is None:
             path = self.path
 
@@ -242,9 +251,13 @@ class TimeSeriesModelBase(ModelBase, ABC):
         return {}
 
     def get_hyperparameters(self) -> dict:
-        """Get hyperparameters that will be passed to the "inner model" that AutoGluon wraps."""
+        """Get dictionary of hyperparameters that will be passed to the "inner model" that AutoGluon wraps."""
         return {**self._get_default_hyperparameters(), **self._hyperparameters}
 
+    def get_hyperparameter(self, key: str) -> Any:
+        """Get a single hyperparameter value for the "inner model"."""
+        return self.get_hyperparameters()[key]
+
     def get_info(self) -> dict:
         """
         Returns a dictionary of numerous fields describing the model.
@@ -384,15 +397,15 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
 
     def __init__(
         self,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        hyperparameters: Optional[dict[str, Any]] = None,
-        freq: Optional[str] = None,
+        path: str | None = None,
+        name: str | None = None,
+        hyperparameters: dict[str, Any] | None = None,
+        freq: str | None = None,
         prediction_length: int = 1,
-        covariate_metadata: Optional[CovariateMetadata] = None,
+        covariate_metadata: CovariateMetadata | None = None,
         target: str = "target",
         quantile_levels: Sequence[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
-        eval_metric: Union[str, TimeSeriesScorer, None] = None,
+        eval_metric: str | TimeSeriesScorer | None = None,
     ):
         # TODO: make freq a required argument in AbstractTimeSeriesModel
         super().__init__(
@@ -406,9 +419,9 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
             quantile_levels=quantile_levels,
             eval_metric=eval_metric,
         )
-        self.target_scaler: Optional[TargetScaler]
-        self.covariate_scaler: Optional[CovariateScaler]
-        self.covariate_regressor: Optional[CovariateRegressor]
+        self.target_scaler: TargetScaler | None
+        self.covariate_scaler: CovariateScaler | None
+        self.covariate_regressor: CovariateRegressor | None
 
     def _initialize_transforms_and_regressor(self) -> None:
         self.target_scaler = get_target_scaler(self.get_hyperparameters().get("target_scaler"), target=self.target)
@@ -433,8 +446,8 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
     def fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
         verbosity: int = 2,
         **kwargs,
     ) -> Self:
@@ -527,10 +540,10 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-        num_cpus: Optional[int] = None,
-        num_gpus: Optional[int] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
         verbosity: int = 2,
         **kwargs,
     ) -> None:
@@ -551,7 +564,7 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
                 "as hyperparameters when initializing or use `hyperparameter_tune` instead."
             )
 
-    def _log_unused_hyperparameters(self, extra_allowed_hyperparameters: Optional[list[str]] = None) -> None:
+    def _log_unused_hyperparameters(self, extra_allowed_hyperparameters: list[str] | None = None) -> None:
         """Log a warning if unused hyperparameters were provided to the model."""
         allowed_hyperparameters = self.allowed_hyperparameters
         if extra_allowed_hyperparameters is not None:
@@ -567,7 +580,7 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
     def predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         """Given a dataset, predict the next `self.prediction_length` time steps.
@@ -648,7 +661,7 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         """Private method for `predict`. See `predict` for documentation of arguments."""
@@ -731,7 +744,7 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
         )
         predict_start_time = time.time()
         oof_predictions = self.predict(past_data, known_covariates=known_covariates, **predict_kwargs)
-        self._oof_predictions = [oof_predictions]
+        self.cache_oof_predictions(oof_predictions)
         if store_predict_time:
             self.predict_time = time.time() - predict_start_time
         if store_val_score:
@@ -740,9 +753,9 @@ class AbstractTimeSeriesModel(TimeSeriesModelBase, TimeSeriesTunable, metaclass=
     def preprocess(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         is_train: bool = False,
         **kwargs,
-    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
         """Method that implements model-specific preprocessing logic."""
         return data, known_covariates

autogluon/timeseries/models/abstract/model_trial.py

@@ -76,7 +76,8 @@ def fit_and_save_model(model, fit_kwargs, train_data, val_data, eval_metric, tim
     time_fit_start = time.time()
     model.fit(train_data=train_data, val_data=val_data, time_limit=time_left, **fit_kwargs)
     model.fit_time = time.time() - time_fit_start
-    model.score_and_cache_oof(val_data, store_val_score=True, store_predict_time=True)
+    if val_data is not None:
+        model.score_and_cache_oof(val_data, store_val_score=True, store_predict_time=True)
 
     logger.debug(f"\tHyperparameter tune run: {model.name}")
     logger.debug(f"\t\t{model.val_score:<7.4f}".ljust(15) + f"= Validation score ({eval_metric.name_with_sign})")

autogluon/timeseries/models/abstract/tunable.py

@@ -5,7 +5,7 @@ import os
 import time
 from abc import ABC, abstractmethod
 from contextlib import nullcontext
-from typing import Any, Optional, Union
+from typing import Any
 
 from typing_extensions import Self
 
@@ -37,12 +37,12 @@ class TimeSeriesTunable(Tunable, ABC):
     def hyperparameter_tune(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame],
+        val_data: TimeSeriesDataFrame | None,
         val_splitter: Any = None,
-        default_num_trials: Optional[int] = 1,
-        refit_every_n_windows: Optional[int] = 1,
-        hyperparameter_tune_kwargs: Union[str, dict] = "auto",
-        time_limit: Optional[float] = None,
+        default_num_trials: int | None = 1,
+        refit_every_n_windows: int | None = 1,
+        hyperparameter_tune_kwargs: str | dict = "auto",
+        time_limit: float | None = None,
     ) -> tuple[dict[str, Any], Any]:
         hpo_executor = self._get_default_hpo_executor()
         hpo_executor.initialize(
@@ -144,13 +144,13 @@ class TimeSeriesTunable(Tunable, ABC):
         """
         return None
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, int | float]:
         return {
             "num_cpus": 1,
         }
 
     def _save_with_data(
-        self, train_data: TimeSeriesDataFrame, val_data: Optional[TimeSeriesDataFrame]
+        self, train_data: TimeSeriesDataFrame, val_data: TimeSeriesDataFrame | None
     ) -> tuple[str, str]:
         self.path = os.path.abspath(self.path)
         self.path_root = self.path.rsplit(self.name, 1)[0]

autogluon/timeseries/models/autogluon_tabular/mlforecast.py

@@ -3,7 +3,7 @@ import logging
 import math
 import time
 import warnings
-from typing import Any, Callable, Collection, Optional, Type, Union
+from typing import Any, Callable, Collection, Type
 
 import numpy as np
 import pandas as pd
@@ -33,7 +33,7 @@ logger = logging.getLogger(__name__)
 class TabularModel(BaseEstimator):
     """A scikit-learn compatible wrapper for arbitrary autogluon.tabular models"""
 
-    def __init__(self, model_class: Type[AbstractTabularModel], model_kwargs: Optional[dict] = None):
+    def __init__(self, model_class: Type[AbstractTabularModel], model_kwargs: dict | None = None):
         self.model_class = model_class
         self.model_kwargs = {} if model_kwargs is None else model_kwargs
         self.feature_pipeline = AutoMLPipelineFeatureGenerator(verbosity=0)
@@ -63,12 +63,12 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
 
     def __init__(
         self,
-        freq: Optional[str] = None,
+        freq: str | None = None,
         prediction_length: int = 1,
-        path: Optional[str] = None,
-        name: Optional[str] = None,
-        eval_metric: Optional[Union[str, TimeSeriesScorer]] = None,
-        hyperparameters: Optional[dict[str, Any]] = None,
+        path: str | None = None,
+        name: str | None = None,
+        eval_metric: str | TimeSeriesScorer | None = None,
+        hyperparameters: dict[str, Any] | None = None,
         **kwargs,
     ):
         super().__init__(
@@ -84,13 +84,13 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         from mlforecast.target_transforms import BaseTargetTransform
 
         self._sum_of_differences: int = 0  # number of time steps removed from each series by differencing
-        self._max_ts_length: Optional[int] = None
+        self._max_ts_length: int | None = None
         self._target_lags: np.ndarray
         self._date_features: list[Callable]
         self._mlf: MLForecast
-        self._scaler: Optional[BaseTargetTransform] = None
+        self._scaler: BaseTargetTransform | None = None
         self._residuals_std_per_item: pd.Series
-        self._train_target_median: Optional[float] = None
+        self._train_target_median: float | None = None
         self._non_boolean_real_covariates: list[str] = []
 
     def _initialize_transforms_and_regressor(self):
@@ -114,10 +114,10 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
     def preprocess(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         is_train: bool = False,
         **kwargs,
-    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
         if is_train:
             # All-NaN series are removed; partially-NaN series in train_data are handled inside _generate_train_val_dfs
             all_nan_items = data.item_ids[
@@ -159,7 +159,11 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         date_features = model_params.get("date_features")
         if date_features is None:
             date_features = get_time_features_for_frequency(self.freq)
-        self._date_features = date_features
+        known_covariates = self.covariate_metadata.known_covariates
+        conflicting = [f.__name__ for f in date_features if f.__name__ in known_covariates]
+        if conflicting:
+            logger.info(f"\tRemoved automatic date_features {conflicting} since they clash with known_covariates")
+        self._date_features = [f for f in date_features if f.__name__ not in known_covariates]
 
         target_transforms = []
         differences = model_params.get("differences")
@@ -212,7 +216,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         return mlforecast_df.groupby(MLF_ITEMID, as_index=False, sort=False).tail(max_length)
 
     def _generate_train_val_dfs(
-        self, data: TimeSeriesDataFrame, max_num_items: Optional[int] = None, max_num_samples: Optional[int] = None
+        self, data: TimeSeriesDataFrame, max_num_items: int | None = None, max_num_samples: int | None = None
     ) -> tuple[pd.DataFrame, pd.DataFrame]:
         # Exclude items that are too short for chosen differences - otherwise exception will be raised
         if self._sum_of_differences > 0:
@@ -266,7 +270,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
     def _to_mlforecast_df(
         self,
         data: TimeSeriesDataFrame,
-        static_features: Optional[pd.DataFrame],
+        static_features: pd.DataFrame | None,
         include_target: bool = True,
     ) -> pd.DataFrame:
         """Convert TimeSeriesDataFrame to a format expected by MLForecast methods `predict` and `preprocess`.
@@ -308,10 +312,10 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
-        num_cpus: Optional[int] = None,
-        num_gpus: Optional[int] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
         verbosity: int = 2,
         **kwargs,
     ) -> None:
@@ -385,8 +389,8 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
     def _remove_short_ts_and_generate_fallback_forecast(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
-    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame], Optional[TimeSeriesDataFrame]]:
+        known_covariates: TimeSeriesDataFrame | None = None,
+    ) -> tuple[TimeSeriesDataFrame, TimeSeriesDataFrame, TimeSeriesDataFrame | None]:
         """Remove series that are too short for chosen differencing from data and generate naive forecast for them.
 
         Returns
@@ -482,7 +486,7 @@ class DirectTabularModel(AbstractMLForecastModel):
     lags : list[int], default = None
         Lags of the target that will be used as features for predictions. If None, will be determined automatically
         based on the frequency of the data.
-    date_features : list[Union[str, Callable]], default = None
+    date_features : list[str | Callable], default = None
         Features computed from the dates. Can be pandas date attributes or functions that will take the dates as input.
         If None, will be determined automatically based on the frequency of the data.
     differences : list[int], default = []
@@ -544,7 +548,7 @@ class DirectTabularModel(AbstractMLForecastModel):
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         from .transforms import apply_inverse_transform
@@ -610,7 +614,7 @@ class DirectTabularModel(AbstractMLForecastModel):
         return predictions_tsdf
 
     def _postprocess_predictions(
-        self, predictions: Union[np.ndarray, pd.Series], repeated_item_ids: pd.Series
+        self, predictions: np.ndarray | pd.Series, repeated_item_ids: pd.Series
     ) -> pd.DataFrame:
         if self.is_quantile_model:
             predictions_df = pd.DataFrame(predictions, columns=[str(q) for q in self.quantile_levels])
@@ -664,7 +668,7 @@ class RecursiveTabularModel(AbstractMLForecastModel):
     lags : list[int], default = None
         Lags of the target that will be used as features for predictions. If None, will be determined automatically
         based on the frequency of the data.
-    date_features : list[Union[str, Callable]], default = None
+    date_features : list[str | Callable], default = None
         Features computed from the dates. Can be pandas date attributes or functions that will take the dates as input.
         If None, will be determined automatically based on the frequency of the data.
     differences : list[int], default = None
@@ -702,7 +706,7 @@ class RecursiveTabularModel(AbstractMLForecastModel):
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         original_item_id_order = data.item_ids

autogluon/timeseries/models/autogluon_tabular/per_step.py

@@ -2,7 +2,7 @@ import logging
 import math
 import os
 import time
-from typing import Any, Callable, Literal, Optional, Type
+from typing import Any, Callable, Literal, Type
 
 import numpy as np
 import pandas as pd
@@ -56,7 +56,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         Seasonal lags of the target used as features. Unlike trailing lags, seasonal lags are not shifted
         but filtered by availability: model for step ``h`` uses ``[lag for lag in seasonal_lags if lag > h]``.
         If None, determined automatically based on data frequency.
-    date_features : list[Union[str, Callable]], default = None
+    date_features : list[str | Callable], default = None
         Features computed from the dates. Can be pandas date attributes or functions that will take the dates as input.
         If None, will be determined automatically based on the frequency of the data.
     target_scaler : {"standard", "mean_abs", "min_max", "robust", None}, default = "mean_abs"
@@ -94,7 +94,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         self._model_cls: Type[AbstractTabularModel]
         self._n_jobs: int
         self._non_boolean_real_covariates: list[str] = []
-        self._max_ts_length: Optional[int] = None
+        self._max_ts_length: int | None = None
 
     @property
     def allowed_hyperparameters(self) -> list[str]:
@@ -140,11 +140,11 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         model_hyperparameters: dict,
         problem_type: Literal["quantile", "regression"],
         eval_metric: str,
-        validation_fraction: Optional[float],
+        validation_fraction: float | None,
         quantile_levels: list[float],
         lags: list[int],
         date_features: list[Callable],
-        time_limit: Optional[float],
+        time_limit: float | None,
         num_cpus: int,
         verbosity: int,
     ) -> str:
@@ -237,7 +237,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
     def preprocess(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         is_train: bool = False,
         **kwargs,
     ):
@@ -263,7 +263,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         return data, known_covariates
 
     def _get_train_df(
-        self, train_data: TimeSeriesDataFrame, max_num_items: Optional[int], max_num_samples: Optional[int]
+        self, train_data: TimeSeriesDataFrame, max_num_items: int | None, max_num_samples: int | None
     ) -> pd.DataFrame:
         if max_num_items is not None and train_data.num_items > max_num_items:
             items_to_keep = train_data.item_ids.to_series().sample(n=int(max_num_items))  # noqa: F841
@@ -305,8 +305,10 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
-        val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[float] = None,
+        val_data: TimeSeriesDataFrame | None = None,
+        time_limit: float | None = None,
+        num_cpus: int | None = None,
+        num_gpus: int | None = None,
         verbosity: int = 2,
         **kwargs,
     ) -> None:
@@ -455,7 +457,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
     def _predict(
         self,
         data: TimeSeriesDataFrame,
-        known_covariates: Optional[TimeSeriesDataFrame] = None,
+        known_covariates: TimeSeriesDataFrame | None = None,
         **kwargs,
     ) -> TimeSeriesDataFrame:
         if known_covariates is not None:

autogluon/timeseries/models/autogluon_tabular/transforms.py

@@ -1,4 +1,4 @@
-from typing import Literal, Union
+from typing import Literal
 
 import numpy as np
 import pandas as pd
@@ -47,7 +47,7 @@ class MLForecastScaler(BaseTargetTransform):
 
 def apply_inverse_transform(
     df: pd.DataFrame,
-    transform: Union[_BaseGroupedArrayTargetTransform, BaseTargetTransform],
+    transform: _BaseGroupedArrayTargetTransform | BaseTargetTransform,
 ) -> pd.DataFrame:
     """Apply inverse transformation to a dataframe, converting to GroupedArray if necessary"""
     if isinstance(transform, BaseTargetTransform):

autogluon/timeseries/models/chronos/__init__.py

@@ -1,3 +1,4 @@
+from .chronos2 import Chronos2Model
 from .model import ChronosModel
 
-__all__ = ["ChronosModel"]
+__all__ = ["ChronosModel", "Chronos2Model"]