autogluon.timeseries 1.4.1b20250821__py3-none-any.whl → 1.4.1b20250823__py3-none-any.whl

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, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Collection, Optional, Type, Union
 
 import numpy as np
 import pandas as pd
@@ -68,7 +68,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         path: Optional[str] = None,
         name: Optional[str] = None,
         eval_metric: Optional[Union[str, TimeSeriesScorer]] = None,
-        hyperparameters: Optional[Dict[str, Any]] = None,
+        hyperparameters: Optional[dict[str, Any]] = None,
         **kwargs,
     ):
         super().__init__(
@@ -86,12 +86,12 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         self._sum_of_differences: int = 0  # number of time steps removed from each series by differencing
         self._max_ts_length: Optional[int] = None
         self._target_lags: np.ndarray
-        self._date_features: List[Callable]
+        self._date_features: list[Callable]
         self._mlf: MLForecast
         self._scaler: Optional[BaseTargetTransform] = None
         self._residuals_std_per_item: pd.Series
         self._train_target_median: Optional[float] = None
-        self._non_boolean_real_covariates: List[str] = []
+        self._non_boolean_real_covariates: list[str] = []
 
     def _initialize_transforms_and_regressor(self):
         super()._initialize_transforms_and_regressor()
@@ -99,7 +99,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         self.target_scaler = None
 
     @property
-    def allowed_hyperparameters(self) -> List[str]:
+    def allowed_hyperparameters(self) -> list[str]:
         return super().allowed_hyperparameters + [
             "lags",
             "date_features",
@@ -117,7 +117,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         is_train: bool = False,
         **kwargs,
-    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
         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[data[self.target].isna().groupby(ITEMID, sort=False).all()]
@@ -130,7 +130,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
                 data[self.target] = data[self.target].fillna(value=self._train_target_median)
         return data, known_covariates
 
-    def _process_deprecated_hyperparameters(self, model_params: Dict[str, Any]) -> Dict[str, Any]:
+    def _process_deprecated_hyperparameters(self, model_params: dict[str, Any]) -> dict[str, Any]:
         if "tabular_hyperparameters" in model_params:
             logger.warning(
                 f"Hyperparameter 'tabular_hyperparameters' for {self.name} is deprecated and will be removed in v1.5. "
@@ -155,7 +155,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             )
         return model_params
 
-    def _get_default_hyperparameters(self) -> Dict[str, Any]:
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
         return {
             "max_num_items": 20_000,
             "max_num_samples": 1_000_000,
@@ -163,12 +163,12 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             "model_hyperparameters": {},
         }
 
-    def _create_tabular_model(self, model_name: str, model_hyperparameters: Dict[str, Any]) -> TabularModel:
+    def _create_tabular_model(self, model_name: str, model_hyperparameters: dict[str, Any]) -> TabularModel:
         raise NotImplementedError
 
     def _get_mlforecast_init_args(
-        self, train_data: TimeSeriesDataFrame, model_params: Dict[str, Any]
-    ) -> Dict[str, Any]:
+        self, train_data: TimeSeriesDataFrame, model_params: dict[str, Any]
+    ) -> dict[str, Any]:
         from mlforecast.target_transforms import Differences
 
         from .transforms import MLForecastScaler
@@ -236,7 +236,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
 
     def _generate_train_val_dfs(
         self, data: TimeSeriesDataFrame, max_num_items: Optional[int] = None, max_num_samples: Optional[int] = None
-    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    ) -> 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:
             ts_lengths = data.num_timesteps_per_item()
@@ -400,16 +400,16 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
         self,
         data: TimeSeriesDataFrame,
         known_covariates: Optional[TimeSeriesDataFrame] = None,
-    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame], Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame], Optional[TimeSeriesDataFrame]]:
         """Remove series that are too short for chosen differencing from data and generate naive forecast for them.
 
         Returns
         -------
-        data_long : TimeSeriesDataFrame
+        data_long
             Data containing only time series that are long enough for the model to predict.
-        known_covariates_long : TimeSeriesDataFrame or None
+        known_covariates_long
             Future known covariates containing only time series that are long enough for the model to predict.
-        forecast_for_short_series : TimeSeriesDataFrame or None
+        forecast_for_short_series
             Seasonal naive forecast for short series, if there are any in the dataset.
         """
         ts_lengths = data.num_timesteps_per_item()
@@ -468,7 +468,7 @@ class AbstractMLForecastModel(AbstractTimeSeriesModel):
             predictions[str(q)] = predictions["mean"] + norm.ppf(q) * std_per_timestep.to_numpy()
         return predictions
 
-    def _more_tags(self) -> Dict[str, Any]:
+    def _more_tags(self) -> dict[str, Any]:
         return {"allow_nan": True, "can_refit_full": True}
 
 
@@ -493,13 +493,13 @@ class DirectTabularModel(AbstractMLForecastModel):
 
     Other Parameters
     ----------------
-    lags : List[int], default = None
+    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[Union[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 = []
+    differences : list[int], default = []
         Differences to take of the target before computing the features. These are restored at the forecasting step.
         Defaults to no differencing.
     target_scaler : {"standard", "mean_abs", "min_max", "robust", None}, default = "mean_abs"
@@ -508,7 +508,7 @@ class DirectTabularModel(AbstractMLForecastModel):
         Name of the tabular regression model. See ``autogluon.tabular.registry.ag_model_registry`` or
         `the documentation <https://auto.gluon.ai/stable/api/autogluon.tabular.models.html>`_ for the list of available
         tabular models.
-    model_hyperparameters : Dict[str, Any], optional
+    model_hyperparameters : dict[str, Any], optional
         Hyperparameters passed to the tabular regression model.
     max_num_items : int or None, default = 20_000
         If not None, the model will randomly select this many time series for training and validation.
@@ -523,7 +523,7 @@ class DirectTabularModel(AbstractMLForecastModel):
     def is_quantile_model(self) -> bool:
         return self.eval_metric.needs_quantile
 
-    def get_hyperparameters(self) -> Dict[str, Any]:
+    def get_hyperparameters(self) -> dict[str, Any]:
         model_params = super().get_hyperparameters()
         # We don't set 'target_scaler' if user already provided 'scaler' to avoid overriding the user-provided value
         if "scaler" not in model_params:
@@ -634,7 +634,7 @@ class DirectTabularModel(AbstractMLForecastModel):
         column_order = ["mean"] + [col for col in predictions_df.columns if col != "mean"]
         return predictions_df[column_order]
 
-    def _create_tabular_model(self, model_name: str, model_hyperparameters: Dict[str, Any]) -> TabularModel:
+    def _create_tabular_model(self, model_name: str, model_hyperparameters: dict[str, Any]) -> TabularModel:
         model_class = ag_model_registry.key_to_cls(model_name)
         if self.is_quantile_model:
             problem_type = ag.constants.QUANTILE
@@ -673,25 +673,25 @@ class RecursiveTabularModel(AbstractMLForecastModel):
 
     Other Parameters
     ----------------
-    lags : List[int], default = None
+    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[Union[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
+    differences : list[int], default = None
         Differences to take of the target before computing the features. These are restored at the forecasting step.
         If None, will be set to ``[seasonal_period]``, where seasonal_period is determined based on the data frequency.
     target_scaler : {"standard", "mean_abs", "min_max", "robust", None}, default = "standard"
         Scaling applied to each time series. Scaling is applied after differencing.
-    lag_transforms : Dict[int, List[Callable]], default = None
+    lag_transforms : dict[int, list[Callable]], default = None
         Dictionary mapping lag periods to transformation functions applied to lagged target values (e.g., rolling mean).
         See `MLForecast documentation <https://nixtlaverse.nixtla.io/mlforecast/lag_transforms.html>`_ for more details.
     model_name : str, default = "GBM"
         Name of the tabular regression model. See ``autogluon.tabular.registry.ag_model_registry`` or
         `the documentation <https://auto.gluon.ai/stable/api/autogluon.tabular.models.html>`_ for the list of available
         tabular models.
-    model_hyperparameters : Dict[str, Any], optional
+    model_hyperparameters : dict[str, Any], optional
         Hyperparameters passed to the tabular regression model.
     max_num_items : int or None, default = 20_000
         If not None, the model will randomly select this many time series for training and validation.
@@ -702,7 +702,7 @@ class RecursiveTabularModel(AbstractMLForecastModel):
 
     ag_priority = 90
 
-    def get_hyperparameters(self) -> Dict[str, Any]:
+    def get_hyperparameters(self) -> dict[str, Any]:
         model_params = super().get_hyperparameters()
         # We don't set 'target_scaler' if user already provided 'scaler' to avoid overriding the user-provided value
         if "scaler" not in model_params:
@@ -756,7 +756,7 @@ class RecursiveTabularModel(AbstractMLForecastModel):
             predictions = pd.concat([predictions, forecast_for_short_series])  # type: ignore
         return predictions.reindex(original_item_id_order, level=ITEMID)
 
-    def _create_tabular_model(self, model_name: str, model_hyperparameters: Dict[str, Any]) -> TabularModel:
+    def _create_tabular_model(self, model_name: str, model_hyperparameters: dict[str, Any]) -> TabularModel:
         model_class = ag_model_registry.key_to_cls(model_name)
         return TabularModel(
             model_class=model_class,

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, Dict, List, Literal, Optional, Type
+from typing import Any, Callable, Literal, Optional, Type
 
 import numpy as np
 import pandas as pd
@@ -49,15 +49,15 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
 
     Other Parameters
     ----------------
-    trailing_lags : List[int], default = None
+    trailing_lags : list[int], default = None
         Trailing window lags of the target that will be used as features for predictions.
         Trailing lags are shifted per forecast step: model for step ``h`` uses ``[lag+h for lag in trailing_lags]``.
         If None, defaults to ``[1, 2, ..., 12]``.
-    seasonal_lags: List[int], default = None
+    seasonal_lags : list[int], default = None
         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[Union[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"
@@ -66,7 +66,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         Name of the tabular regression model. See ``autogluon.tabular.registry.ag_model_registry`` or
         `the documentation <https://auto.gluon.ai/stable/api/autogluon.tabular.models.html>`_ for the list of available
         tabular models.
-    model_hyperparameters : Dict[str, Any], optional
+    model_hyperparameters : dict[str, Any], optional
         Hyperparameters passed to the tabular regression model.
     validation_fraction : float or None, default = 0.1
         Fraction of the training data to use for validation. If None or 0.0, no validation set is created.
@@ -94,11 +94,11 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         self._date_features: list[Callable]
         self._model_cls: Type[AbstractTabularModel]
         self._n_jobs: int
-        self._non_boolean_real_covariates: List[str] = []
+        self._non_boolean_real_covariates: list[str] = []
         self._max_ts_length: Optional[int] = None
 
     @property
-    def allowed_hyperparameters(self) -> List[str]:
+    def allowed_hyperparameters(self) -> list[str]:
         # TODO: Differencing is currently not supported because it greatly complicates the preprocessing logic
         return super().allowed_hyperparameters + [
             "trailing_lags",
@@ -285,10 +285,10 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
 
     @staticmethod
     def _get_lags_for_step(
-        trailing_lags: List[int],
-        seasonal_lags: List[int],
+        trailing_lags: list[int],
+        seasonal_lags: list[int],
         step: int,
-    ) -> List[int]:
+    ) -> list[int]:
         """Get the list of lags that can be used by the model for the given step."""
         shifted_trailing_lags = [lag + step for lag in trailing_lags]
         # Only keep lags that are available for model predicting `step` values ahead at prediction time
@@ -417,7 +417,7 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
 
         Returns
         -------
-        predictions :
+        predictions
             Predictions of the model for the given step. Shape: (num_items, len(quantile_levels)).
         """
         from mlforecast import MLForecast
@@ -498,5 +498,5 @@ class PerStepTabularModel(AbstractTimeSeriesModel):
         predictions["mean"] = predictions["0.5"]
         return TimeSeriesDataFrame(predictions)
 
-    def _more_tags(self) -> Dict[str, Any]:
+    def _more_tags(self) -> dict[str, Any]:
         return {"allow_nan": True, "can_refit_full": True}

autogluon/timeseries/models/chronos/model.py

@@ -3,7 +3,7 @@ import os
 import shutil
 import warnings
 from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -110,7 +110,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
     Other Parameters
     ----------------
-    model_path: str, default = "autogluon/chronos-bolt-small"
+    model_path : str, default = "autogluon/chronos-bolt-small"
         Model path used for the model, i.e., a HuggingFace transformers ``name_or_path``. Can be a
         compatible model name on HuggingFace Hub or a local path to a model directory. Original
         Chronos models (i.e., ``autogluon/chronos-t5-{model_size}``) can be specified with aliases
@@ -144,7 +144,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         for more information.
     fine_tune : bool, default = False
         If True, the pretrained model will be fine-tuned
-    fine_tune_lr: float, default = 1e-5
+    fine_tune_lr : float, default = 1e-5
         The learning rate used for fine-tuning. This default is suitable for Chronos-Bolt models; for the original
         Chronos models, we recommend using a higher learning rate such as ``1e-4``
     fine_tune_steps : int, default = 1000
@@ -162,7 +162,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         during fine-tuning. If None, the entire validation dataset will be used.
     fine_tune_trainer_kwargs : dict, optional
         Extra keyword arguments passed to ``transformers.TrainingArguments``
-    keep_transformers_logs: bool, default = False
+    keep_transformers_logs : bool, default = False
         If True, the logs generated by transformers will NOT be removed after fine-tuning
     """
 
@@ -181,7 +181,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         path: Optional[str] = None,
         name: Optional[str] = None,
         eval_metric: Optional[str] = None,
-        hyperparameters: Optional[Dict[str, Any]] = None,
+        hyperparameters: Optional[dict[str, Any]] = None,
         **kwargs,  # noqa
     ):
         hyperparameters = hyperparameters if hyperparameters is not None else {}
@@ -242,7 +242,7 @@ class ChronosModel(AbstractTimeSeriesModel):
         return self._model_pipeline
 
     @property
-    def ag_default_config(self) -> Dict[str, Any]:
+    def ag_default_config(self) -> dict[str, Any]:
         """The default configuration of the model used by AutoGluon if the model is one of those
         defined in MODEL_CONFIGS. For now, these are ``autogluon/chronos-t5-*`` family of models.
         """
@@ -272,8 +272,8 @@ class ChronosModel(AbstractTimeSeriesModel):
         """
         return self.ag_default_config.get("default_torch_dtype", "auto")
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> Dict[str, Union[int, float]]:
-        minimum_resources: Dict[str, Union[int, float]] = {"num_cpus": 1}
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
+        minimum_resources: dict[str, Union[int, float]] = {"num_cpus": 1}
         # if GPU is available, we train with 1 GPU per trial
         if is_gpu_available:
             minimum_resources["num_gpus"] = self.min_num_gpus
@@ -323,7 +323,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
         return init_args.copy()
 
-    def _get_default_hyperparameters(self) -> Dict:
+    def _get_default_hyperparameters(self) -> dict:
         return {
             "batch_size": self.default_batch_size,
             "num_samples": self.default_num_samples,
@@ -688,7 +688,7 @@ class ChronosModel(AbstractTimeSeriesModel):
 
         return TimeSeriesDataFrame(df)
 
-    def _more_tags(self) -> Dict:
+    def _more_tags(self) -> dict:
         do_fine_tune = self.get_hyperparameters()["fine_tune"]
         return {
             "allow_nan": True,

autogluon/timeseries/models/chronos/pipeline/base.py

@@ -2,7 +2,7 @@
 
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import torch
 
@@ -18,7 +18,7 @@ class ForecastType(Enum):
 
 
 class PipelineRegistry(type):
-    REGISTRY: Dict[str, "PipelineRegistry"] = {}
+    REGISTRY: dict[str, "PipelineRegistry"] = {}
 
     def __new__(cls, name, bases, attrs):
         """See, https://github.com/faif/python-patterns."""
@@ -43,13 +43,13 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
         """
         Parameters
         ----------
-        inner_model : PreTrainedModel
+        inner_model
             A hugging-face transformers PreTrainedModel, e.g., T5ForConditionalGeneration
         """
         # for easy access to the inner HF-style model
         self.inner_model = inner_model
 
-    def _prepare_and_validate_context(self, context: Union[torch.Tensor, List[torch.Tensor]]):
+    def _prepare_and_validate_context(self, context: Union[torch.Tensor, list[torch.Tensor]]):
         if isinstance(context, list):
             context = left_pad_and_stack_1D(context)
         assert isinstance(context, torch.Tensor)
@@ -61,7 +61,7 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
 
     def predict(
         self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
+        context: Union[torch.Tensor, list[torch.Tensor]],
         prediction_length: Optional[int] = None,
         **kwargs,
     ):
@@ -88,8 +88,8 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
         raise NotImplementedError()
 
     def predict_quantiles(
-        self, context: torch.Tensor, prediction_length: int, quantile_levels: List[float], **kwargs
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, context: torch.Tensor, prediction_length: int, quantile_levels: list[float], **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Get quantile and mean forecasts for given time series. All
         predictions are returned on the CPU.
@@ -104,7 +104,7 @@ class BaseChronosPipeline(metaclass=PipelineRegistry):
         prediction_length
             Time steps to predict. Defaults to a model-dependent
             value if not given.
-        quantile_levels: List[float]
+        quantile_levels
             Quantile levels to compute
 
         Returns

autogluon/timeseries/models/chronos/pipeline/chronos.py

@@ -7,7 +7,7 @@
 import logging
 import warnings
 from dataclasses import dataclass
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Literal, Optional, Union
 
 import torch
 import torch.nn as nn
@@ -29,7 +29,7 @@ class ChronosConfig:
     """
 
     tokenizer_class: str
-    tokenizer_kwargs: Dict[str, Any]
+    tokenizer_kwargs: dict[str, Any]
     n_tokens: int
     n_special_tokens: int
     pad_token_id: int
@@ -66,7 +66,7 @@ class ChronosTokenizer:
     def context_input_transform(
         self,
         context: torch.Tensor,
-    ) -> Tuple:
+    ) -> tuple:
         """
         Turn a batch of time series into token IDs, attention mask, and tokenizer_state.
 
@@ -95,7 +95,7 @@ class ChronosTokenizer:
         """
         raise NotImplementedError()
 
-    def label_input_transform(self, label: torch.Tensor, tokenizer_state: Any) -> Tuple:
+    def label_input_transform(self, label: torch.Tensor, tokenizer_state: Any) -> tuple:
         """
         Turn a batch of label slices of time series into token IDs and attention mask
         using the ``tokenizer_state`` provided by ``context_input_transform``.
@@ -171,7 +171,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
     def _input_transform(
         self, context: torch.Tensor, scale: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         attention_mask = ~torch.isnan(context)
 
         if scale is None:
@@ -196,7 +196,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
     def _append_eos_token(
         self, token_ids: torch.Tensor, attention_mask: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         batch_size = token_ids.shape[0]
         eos_tokens = torch.full((batch_size, 1), fill_value=self.config.eos_token_id)
         token_ids = torch.concat((token_ids, eos_tokens), dim=1)
@@ -205,7 +205,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
         return token_ids, attention_mask
 
-    def context_input_transform(self, context: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def context_input_transform(self, context: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         length = context.shape[-1]
 
         if length > self.config.context_length:
@@ -218,7 +218,7 @@ class MeanScaleUniformBins(ChronosTokenizer):
 
         return token_ids, attention_mask, scale
 
-    def label_input_transform(self, label: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    def label_input_transform(self, label: torch.Tensor, scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         token_ids, attention_mask, _ = self._input_transform(context=label, scale=scale)
 
         if self.config.use_eos_token:
@@ -371,7 +371,7 @@ class ChronosPipeline(BaseChronosPipeline):
         self.model = model
 
     @torch.no_grad()
-    def embed(self, context: Union[torch.Tensor, List[torch.Tensor]]) -> Tuple[torch.Tensor, Any]:
+    def embed(self, context: Union[torch.Tensor, list[torch.Tensor]]) -> tuple[torch.Tensor, Any]:
         """
         Get encoder embeddings for the given time series.
 
@@ -404,7 +404,7 @@ class ChronosPipeline(BaseChronosPipeline):
 
     def predict(
         self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
+        context: Union[torch.Tensor, list[torch.Tensor]],
         prediction_length: Optional[int] = None,
         num_samples: Optional[int] = None,
         temperature: Optional[float] = None,
@@ -494,10 +494,10 @@ class ChronosPipeline(BaseChronosPipeline):
         self,
         context: torch.Tensor,
         prediction_length: int,
-        quantile_levels: List[float],
+        quantile_levels: list[float],
         num_samples: Optional[int] = None,
         **kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         num_samples = num_samples or self.model.config.num_samples
         prediction_samples = (
             self.predict(

autogluon/timeseries/models/chronos/pipeline/chronos_bolt.py

@@ -7,7 +7,7 @@ import copy
 import logging
 import warnings
 from dataclasses import dataclass, fields
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -32,7 +32,7 @@ class ChronosBoltConfig:
     prediction_length: int
     input_patch_size: int
     input_patch_stride: int
-    quantiles: List[float]
+    quantiles: list[float]
     use_reg_token: bool = False
 
 
@@ -77,8 +77,8 @@ class InstanceNorm(nn.Module):
     def forward(
         self,
         x: torch.Tensor,
-        loc_scale: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        loc_scale: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         if loc_scale is None:
             loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0)
             scale = torch.nan_to_num((x - loc).square().nanmean(dim=-1, keepdim=True).sqrt(), nan=1.0)
@@ -88,7 +88,7 @@ class InstanceNorm(nn.Module):
 
         return (x - loc) / scale, (loc, scale)
 
-    def inverse(self, x: torch.Tensor, loc_scale: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
+    def inverse(self, x: torch.Tensor, loc_scale: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
         loc, scale = loc_scale
         return x * scale + loc
 
@@ -343,11 +343,11 @@ class ChronosBoltModelForForecasting(T5PreTrainedModel):
         """
         Parameters
         ----------
-        input_embeds: torch.Tensor
+        input_embeds
             Patched and embedded inputs. Shape (batch_size, patched_context_length, d_model)
-        attention_mask: torch.Tensor
+        attention_mask
             Attention mask for the patched context. Shape (batch_size, patched_context_length), type: torch.int64
-        hidden_states: torch.Tensor
+        hidden_states
             Hidden states returned by the encoder. Shape (batch_size, patched_context_length, d_model)
 
         Returns
@@ -385,12 +385,12 @@ class ChronosBoltPipeline(BaseChronosPipeline):
         self.model_prediction_length: int = self.model.config.chronos_config["prediction_length"]
 
     @property
-    def quantiles(self) -> List[float]:
+    def quantiles(self) -> list[float]:
         return self.model.config.chronos_config["quantiles"]
 
     def predict(  # type: ignore[override]
         self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
+        context: Union[torch.Tensor, list[torch.Tensor]],
         prediction_length: Optional[int] = None,
         limit_prediction_length: bool = False,
     ):
@@ -458,8 +458,8 @@ class ChronosBoltPipeline(BaseChronosPipeline):
         return torch.cat(predictions, dim=-1)[..., :prediction_length]
 
     def predict_quantiles(
-        self, context: torch.Tensor, prediction_length: int, quantile_levels: List[float], **kwargs
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, context: torch.Tensor, prediction_length: int, quantile_levels: list[float], **kwargs
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         # shape (batch_size, prediction_length, len(training_quantile_levels))
         predictions = (
             self.predict(