autogluon.timeseries 1.3.2b20250712__py3-none-any.whl → 1.4.1b20251116__py3-none-any.whl

autogluon/timeseries/models/local/abstract_local_model.py

@@ -1,7 +1,7 @@
 import logging
 import time
 from multiprocessing import TimeoutError
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -9,7 +9,7 @@ from joblib import Parallel, cpu_count, delayed
 from scipy.stats import norm
 
 from autogluon.core.utils.exceptions import TimeLimitExceeded
-from autogluon.timeseries.dataset.ts_dataframe import ITEMID, TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.metrics import TimeSeriesScorer
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.utils.datetime import get_seasonality
@@ -29,17 +29,17 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
 
     Attributes
     ----------
-    allowed_local_model_args : List[str]
+    allowed_local_model_args
         Argument that can be passed to the underlying local model.
-    default_max_ts_length : Optional[int]
+    default_max_ts_length
         If not None, only the last ``max_ts_length`` time steps of each time series will be used to train the model.
         This significantly speeds up fitting and usually leads to no change in accuracy.
-    init_time_in_seconds : int
+    init_time_in_seconds
         Time that it takes to initialize the model in seconds (e.g., because of JIT compilation by Numba).
         If time_limit is below this number, model won't be trained.
     """
 
-    allowed_local_model_args: List[str] = []
+    allowed_local_model_args: list[str] = []
     default_max_ts_length: Optional[int] = 2500
     default_max_time_limit_ratio = 1.0
     init_time_in_seconds: int = 0
@@ -51,7 +51,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         path: Optional[str] = None,
         name: Optional[str] = None,
         eval_metric: Union[str, TimeSeriesScorer, None] = None,
-        hyperparameters: Optional[Dict[str, Any]] = None,
+        hyperparameters: Optional[dict[str, Any]] = None,
         **kwargs,  # noqa
     ):
         super().__init__(
@@ -64,12 +64,12 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
             **kwargs,
         )
 
-        self._local_model_args: Dict[str, Any]
+        self._local_model_args: dict[str, Any]
         self._seasonal_period: int
         self._dummy_forecast: pd.DataFrame
 
     @property
-    def allowed_hyperparameters(self) -> List[str]:
+    def allowed_hyperparameters(self) -> list[str]:
         return (
             super().allowed_hyperparameters
             + ["use_fallback_model", "max_ts_length", "n_jobs"]
@@ -82,7 +82,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         known_covariates: Optional[TimeSeriesDataFrame] = None,
         is_train: bool = False,
         **kwargs,
-    ) -> Tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
+    ) -> tuple[TimeSeriesDataFrame, Optional[TimeSeriesDataFrame]]:
         if not self._get_tags()["allow_nan"]:
             data = data.fill_missing_values()
         return data, known_covariates
@@ -134,7 +134,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         stats_repeated = np.tile(stats_marginal.values, [self.prediction_length, 1])
         return pd.DataFrame(stats_repeated, columns=stats_marginal.index)
 
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         return local_model_args
 
     def _predict(self, data: TimeSeriesDataFrame, **kwargs) -> TimeSeriesDataFrame:
@@ -145,7 +145,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
             data = data.slice_by_timestep(-max_ts_length, None)
 
         indptr = data.get_indptr()
-        target_series = data[self.target].droplevel(level=ITEMID)
+        target_series = data[self.target].droplevel(level=TimeSeriesDataFrame.ITEMID)
         all_series = (target_series[indptr[i] : indptr[i + 1]] for i in range(len(indptr) - 1))
 
         # timeout ensures that no individual job takes longer than time_limit
@@ -185,7 +185,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
         time_series: pd.Series,
         use_fallback_model: bool,
         end_time: Optional[float] = None,
-    ) -> Tuple[pd.DataFrame, bool]:
+    ) -> tuple[pd.DataFrame, bool]:
         if end_time is not None and time.time() >= end_time:
             raise TimeLimitExceeded
 
@@ -222,7 +222,7 @@ class AbstractLocalModel(AbstractTimeSeriesModel):
 
 
 def seasonal_naive_forecast(
-    target: np.ndarray, prediction_length: int, quantile_levels: List[float], seasonal_period: int
+    target: np.ndarray, prediction_length: int, quantile_levels: list[float], seasonal_period: int
 ) -> pd.DataFrame:
     """Generate seasonal naive forecast, predicting the last observed value from the same period."""
 

autogluon/timeseries/models/local/naive.py

@@ -24,6 +24,7 @@ class NaiveModel(AbstractLocalModel):
         When set to -1, all CPU cores are used.
     """
 
+    ag_priority = 100
     allowed_local_model_args = ["seasonal_period"]
 
     def _predict_with_local_model(
@@ -66,6 +67,7 @@ class SeasonalNaiveModel(AbstractLocalModel):
         When set to -1, all CPU cores are used.
     """
 
+    ag_priority = 100
     allowed_local_model_args = ["seasonal_period"]
 
     def _predict_with_local_model(
@@ -99,6 +101,7 @@ class AverageModel(AbstractLocalModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 100
     allowed_local_model_args = ["seasonal_period"]
     default_max_ts_length = None
 
@@ -138,6 +141,7 @@ class SeasonalAverageModel(AbstractLocalModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 100
     allowed_local_model_args = ["seasonal_period"]
     default_max_ts_length = None
 

autogluon/timeseries/models/local/npts.py

@@ -36,6 +36,7 @@ class NPTSModel(AbstractLocalModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 80
     allowed_local_model_args = [
         "kernel_type",
         "exp_kernel_weights",

autogluon/timeseries/models/local/statsforecast.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Dict, Optional, Type
+from typing import Any, Optional, Type
 
 import numpy as np
 import pandas as pd
@@ -14,7 +14,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
 
     init_time_in_seconds = 15  # numba compilation for the first run
 
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         seasonal_period = local_model_args.pop("seasonal_period")
         local_model_args["season_length"] = seasonal_period
         return local_model_args
@@ -22,7 +22,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
     def _get_model_type(self, variant: Optional[str] = None) -> Type:
         raise NotImplementedError
 
-    def _get_local_model(self, local_model_args: Dict):
+    def _get_local_model(self, local_model_args: dict):
         local_model_args = local_model_args.copy()
         variant = local_model_args.pop("variant", None)
         model_type = self._get_model_type(variant)
@@ -31,7 +31,7 @@ class AbstractStatsForecastModel(AbstractLocalModel):
     def _get_point_forecast(
         self,
         time_series: pd.Series,
-        local_model_args: Dict,
+        local_model_args: dict,
     ) -> np.ndarray:
         return self._get_local_model(local_model_args).forecast(
             h=self.prediction_length, y=time_series.values.ravel()
@@ -133,6 +133,7 @@ class AutoARIMAModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 60
     init_time_in_seconds = 0  # C++ models require no compilation
     allowed_local_model_args = [
         "d",
@@ -175,9 +176,9 @@ class ARIMAModel(AbstractProbabilisticStatsForecastModel):
 
     Other Parameters
     ----------------
-    order: Tuple[int, int, int], default = (1, 1, 1)
+    order: tuple[int, int, int], default = (1, 1, 1)
         The (p, d, q) order of the model for the number of AR parameters, differences, and MA parameters to use.
-    seasonal_order: Tuple[int, int, int], default = (0, 0, 0)
+    seasonal_order: tuple[int, int, int], default = (0, 0, 0)
         The (P, D, Q) parameters of the seasonal ARIMA model. Setting to (0, 0, 0) disables seasonality.
     include_mean : bool, default = True
         Should the ARIMA model include a mean term?
@@ -193,7 +194,7 @@ class ARIMAModel(AbstractProbabilisticStatsForecastModel):
     method : {"CSS-ML", "CSS", "ML"}, default = "CSS-ML"
         Fitting method: CSS (conditional sum of squares), ML (maximum likelihood), CSS-ML (initialize with CSS, then
         optimize with ML).
-    fixed : Dict[str, float], optional
+    fixed : dict[str, float], optional
         Dictionary containing fixed coefficients for the ARIMA model.
     seasonal_period : int or None, default = None
         Number of time steps in a complete seasonal cycle for seasonal models. For example, 7 for daily data with a
@@ -211,6 +212,7 @@ class ARIMAModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 10
     init_time_in_seconds = 0  # C++ models require no compilation
     allowed_local_model_args = [
         "order",
@@ -267,6 +269,7 @@ class AutoETSModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 70
     init_time_in_seconds = 0  # C++ models require no compilation
     allowed_local_model_args = [
         "damped",
@@ -330,6 +333,8 @@ class ETSModel(AutoETSModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 80
+
     def _update_local_model_args(self, local_model_args: dict) -> dict:
         local_model_args = super()._update_local_model_args(local_model_args)
         local_model_args.setdefault("model", "AAA")
@@ -369,6 +374,7 @@ class DynamicOptimizedThetaModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 75
     allowed_local_model_args = [
         "decomposition_type",
         "seasonal_period",
@@ -413,6 +419,7 @@ class ThetaModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 75
     allowed_local_model_args = [
         "decomposition_type",
         "seasonal_period",
@@ -442,7 +449,7 @@ class AbstractConformalizedStatsForecastModel(AbstractStatsForecastModel):
     def _get_nonconformity_scores(
         self,
         time_series: pd.Series,
-        local_model_args: Dict,
+        local_model_args: dict,
     ) -> np.ndarray:
         h = self.prediction_length
         y = time_series.values.ravel()
@@ -533,6 +540,7 @@ class AutoCESModel(AbstractProbabilisticStatsForecastModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 10
     allowed_local_model_args = [
         "model",
         "seasonal_period",
@@ -548,7 +556,7 @@ class AutoCESModel(AbstractProbabilisticStatsForecastModel):
         local_model_args.setdefault("model", "Z")
         return local_model_args
 
-    def _get_point_forecast(self, time_series: pd.Series, local_model_args: Dict):
+    def _get_point_forecast(self, time_series: pd.Series, local_model_args: dict):
         # Disable seasonality if time series too short for chosen season_length or season_length == 1,
         # otherwise model will crash
         if len(time_series) < 5:
@@ -560,7 +568,7 @@ class AutoCESModel(AbstractProbabilisticStatsForecastModel):
 
 
 class AbstractStatsForecastIntermittentDemandModel(AbstractConformalizedStatsForecastModel):
-    def _update_local_model_args(self, local_model_args: Dict[str, Any]) -> Dict[str, Any]:
+    def _update_local_model_args(self, local_model_args: dict[str, Any]) -> dict[str, Any]:
         _ = local_model_args.pop("seasonal_period")
         return local_model_args
 
@@ -600,6 +608,8 @@ class ADIDAModel(AbstractStatsForecastIntermittentDemandModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 10
+
     def _get_model_type(self, variant: Optional[str] = None):
         from statsforecast.models import ADIDA
 
@@ -622,9 +632,9 @@ class CrostonModel(AbstractStatsForecastIntermittentDemandModel):
     variant : {"SBA", "classic", "optimized"}, default = "SBA"
         Variant of the Croston model that is used. Available options:
 
-        - `"classic"` - variant of the Croston method where the smoothing parameter is fixed to 0.1 (based on `statsforecast.models.CrostonClassic <https://nixtla.mintlify.app/statsforecast/docs/models/crostonclassic.html>`_)
-        - `"SBA"` - variant of the Croston method based on Syntetos-Boylan Approximation (based on `statsforecast.models.CrostonSBA <https://nixtla.mintlify.app/statsforecast/docs/models/crostonsba.html>`_)
-        - `"optimized"` - variant of the Croston method where the smoothing parameter is optimized (based on `statsforecast.models.CrostonOptimized <https://nixtla.mintlify.app/statsforecast/docs/models/crostonoptimized.html>`_)
+        - ``"classic"`` - variant of the Croston method where the smoothing parameter is fixed to 0.1 (based on `statsforecast.models.CrostonClassic <https://nixtla.mintlify.app/statsforecast/docs/models/crostonclassic.html>`_)
+        - ``"SBA"`` - variant of the Croston method based on Syntetos-Boylan Approximation (based on `statsforecast.models.CrostonSBA <https://nixtla.mintlify.app/statsforecast/docs/models/crostonsba.html>`_)
+        - ``"optimized"`` - variant of the Croston method where the smoothing parameter is optimized (based on `statsforecast.models.CrostonOptimized <https://nixtla.mintlify.app/statsforecast/docs/models/crostonoptimized.html>`_)
 
     n_jobs : int or float, default = joblib.cpu_count(only_physical_cores=True)
         Number of CPU cores used to fit the models in parallel.
@@ -636,6 +646,8 @@ class CrostonModel(AbstractStatsForecastIntermittentDemandModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_model_aliases = ["CrostonSBA"]
+    ag_priority = 80
     allowed_local_model_args = [
         "variant",
     ]
@@ -688,6 +700,8 @@ class IMAPAModel(AbstractStatsForecastIntermittentDemandModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 10
+
     def _get_model_type(self, variant: Optional[str] = None):
         from statsforecast.models import IMAPA
 
@@ -710,6 +724,8 @@ class ZeroModel(AbstractStatsForecastIntermittentDemandModel):
         This significantly speeds up fitting and usually leads to no change in accuracy.
     """
 
+    ag_priority = 100
+
     def _get_model_type(self, variant: Optional[str] = None):
         # ZeroModel does not depend on a StatsForecast implementation
         raise NotImplementedError
@@ -717,6 +733,6 @@ class ZeroModel(AbstractStatsForecastIntermittentDemandModel):
     def _get_point_forecast(
         self,
         time_series: pd.Series,
-        local_model_args: Dict,
+        local_model_args: dict,
     ):
         return np.zeros(self.prediction_length)

autogluon/timeseries/models/multi_window/multi_window_model.py

@@ -4,12 +4,13 @@ import logging
 import math
 import os
 import time
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Optional, Type, Union
 
 import numpy as np
+from typing_extensions import Self
 
 import autogluon.core as ag
-from autogluon.timeseries.dataset.ts_dataframe import TimeSeriesDataFrame
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.models.abstract import AbstractTimeSeriesModel
 from autogluon.timeseries.models.local.abstract_local_model import AbstractLocalModel
 from autogluon.timeseries.splitter import AbstractWindowSplitter, ExpandingWindowSplitter
@@ -25,10 +26,10 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
     Parameters
     ----------
-    model_base : Union[AbstractTimeSeriesModel, Type[AbstractTimeSeriesModel]]
+    model_base
         The base model to repeatedly train. If a AbstractTimeSeriesModel class, then also provide model_base_kwargs
         which will be used to initialize the model via model_base(**model_base_kwargs).
-    model_base_kwargs : Optional[Dict[str, any]], default = None
+    model_base_kwargs
         kwargs used to initialize model_base if model_base is a class.
     """
 
@@ -38,7 +39,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def __init__(
         self,
         model_base: Union[AbstractTimeSeriesModel, Type[AbstractTimeSeriesModel]],
-        model_base_kwargs: Optional[Dict[str, Any]] = None,
+        model_base_kwargs: Optional[dict[str, Any]] = None,
         **kwargs,
     ):
         if inspect.isclass(model_base) and issubclass(model_base, AbstractTimeSeriesModel):
@@ -73,10 +74,6 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def supports_past_covariates(self) -> bool:
         return self.model_base.supports_past_covariates
 
-    @property
-    def supports_cat_covariates(self) -> bool:
-        return self.model_base.supports_cat_covariates
-
     def _get_model_base(self):
         return self.model_base
 
@@ -86,15 +83,18 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
     def _is_gpu_available(self) -> bool:
         return self._get_model_base()._is_gpu_available()
 
-    def get_minimum_resources(self, is_gpu_available: bool = False) -> bool:
+    def get_minimum_resources(self, is_gpu_available: bool = False) -> dict[str, Union[int, float]]:
         return self._get_model_base().get_minimum_resources(is_gpu_available)
 
     def _fit(
         self,
         train_data: TimeSeriesDataFrame,
         val_data: Optional[TimeSeriesDataFrame] = None,
-        time_limit: Optional[int] = None,
-        val_splitter: AbstractWindowSplitter = None,
+        time_limit: Optional[float] = None,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
+        verbosity: int = 2,
+        val_splitter: Optional[AbstractWindowSplitter] = None,
         refit_every_n_windows: Optional[int] = 1,
         **kwargs,
     ):
@@ -111,11 +111,15 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
         oof_predictions_per_window = []
         global_fit_start_time = time.time()
+        model: Optional[AbstractTimeSeriesModel] = None
 
         for window_index, (train_fold, val_fold) in enumerate(val_splitter.split(train_data)):
             logger.debug(f"\tWindow {window_index}")
+
             # refit_this_window is always True for the 0th window
             refit_this_window = window_index % refit_every_n_windows == 0
+            assert window_index != 0 or refit_this_window
+
             if time_limit is None:
                 time_left_for_window = None
             else:
@@ -148,6 +152,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
             else:
                 time_left_for_prediction = time_limit - (time.time() - global_fit_start_time)
 
+            assert model is not None
             model.score_and_cache_oof(
                 val_fold, store_val_score=True, store_predict_time=True, time_limit=time_left_for_prediction
             )
@@ -172,11 +177,13 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
 
         # Only the model trained on most recent data is saved & used for prediction
         self.most_recent_model = model
-        self.most_recent_model_folder = most_recent_refit_window
+        assert self.most_recent_model is not None
+
+        self.most_recent_model_folder = most_recent_refit_window  # type: ignore
         self.predict_time = self.most_recent_model.predict_time
-        self.fit_time = time.time() - global_fit_start_time - self.predict_time
+        self.fit_time = time.time() - global_fit_start_time - self.predict_time  # type: ignore
         self._oof_predictions = oof_predictions_per_window
-        self.val_score = np.mean([info["val_score"] for info in self.info_per_val_window])
+        self.val_score = np.mean([info["val_score"] for info in self.info_per_val_window])  # type: ignore
 
     def get_info(self) -> dict:
         info = super().get_info()
@@ -227,7 +234,7 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
         train_fn_kwargs["init_params"]["model_base_kwargs"] = self.get_params()
         return train_fn_kwargs
 
-    def save(self, path: str = None, verbose=True) -> str:
+    def save(self, path: Optional[str] = None, verbose: bool = True) -> str:
         most_recent_model = self.most_recent_model
         self.most_recent_model = None
         save_path = super().save(path, verbose)
@@ -238,32 +245,36 @@ class MultiWindowBacktestingModel(AbstractTimeSeriesModel):
             most_recent_model.save()
         return save_path
 
-    def persist(self):
+    def persist(self) -> Self:
         if self.most_recent_model is None:
             raise ValueError(f"{self.name} must be fit before persisting")
         self.most_recent_model.persist()
+        return self
 
     @classmethod
     def load(
         cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True
     ) -> AbstractTimeSeriesModel:
         model = super().load(path=path, reset_paths=reset_paths, load_oof=load_oof, verbose=verbose)
-        most_recent_model_path = os.path.join(model.path, model.most_recent_model_folder)
-        model.most_recent_model = model.model_base_type.load(
-            most_recent_model_path,
-            reset_paths=reset_paths,
-            verbose=verbose,
-        )
+        if model.most_recent_model_folder is not None:
+            most_recent_model_path = os.path.join(model.path, model.most_recent_model_folder)
+            model.most_recent_model = model.model_base_type.load(
+                most_recent_model_path,
+                reset_paths=reset_paths,
+                verbose=verbose,
+            )
         return model
 
     def convert_to_refit_full_template(self) -> AbstractTimeSeriesModel:
         # refit_model is an instance of base model type, not MultiWindowBacktestingModel
+        assert self.most_recent_model is not None, "Most recent model is None. Model must be fit first."
         refit_model = self.most_recent_model.convert_to_refit_full_template()
         refit_model.rename(self.name + ag.constants.REFIT_FULL_SUFFIX)
         return refit_model
 
     def convert_to_refit_full_via_copy(self) -> AbstractTimeSeriesModel:
         # refit_model is an instance of base model type, not MultiWindowBacktestingModel
+        assert self.most_recent_model is not None, "Most recent model is None. Model must be fit first."
         refit_model = self.most_recent_model.convert_to_refit_full_via_copy()
         refit_model.rename(self.name + ag.constants.REFIT_FULL_SUFFIX)
         return refit_model

autogluon/timeseries/models/registry.py

@@ -0,0 +1,65 @@
+from abc import ABCMeta
+from dataclasses import dataclass
+from inspect import isabstract
+from typing import Union
+
+
+@dataclass
+class ModelRecord:
+    model_class: type
+    ag_priority: int
+
+
+class ModelRegistry(ABCMeta):
+    """Registry metaclass for time series models. Ensures that TimeSeriesModel classes
+    which implement this metaclass are automatically registered, in order to centralize
+    access to model types.
+
+    See, https://github.com/faif/python-patterns.
+    """
+
+    REGISTRY: dict[str, ModelRecord] = {}
+
+    def __new__(cls, name, bases, attrs):
+        new_cls = super().__new__(cls, name, bases, attrs)
+
+        if name is not None and not isabstract(new_cls):
+            record = ModelRecord(
+                model_class=new_cls,
+                ag_priority=getattr(new_cls, "ag_priority", 0),
+            )
+            cls._add(name.removesuffix("Model"), record)
+
+            # if the class provides additional aliases, register them too
+            if aliases := attrs.get("ag_model_aliases"):
+                for alias in aliases:
+                    cls._add(alias, record)
+
+        return new_cls
+
+    @classmethod
+    def _add(cls, alias: str, record: ModelRecord) -> None:
+        if alias in cls.REGISTRY:
+            raise ValueError(f"You are trying to define a new model with {alias}, but this model already exists.")
+        cls.REGISTRY[alias] = record
+
+    @classmethod
+    def _get_model_record(cls, alias: Union[str, type]) -> ModelRecord:
+        if isinstance(alias, type):
+            alias = alias.__name__
+        alias = alias.removesuffix("Model")
+        if alias not in cls.REGISTRY:
+            raise ValueError(f"Unknown model: {alias}, available models are: {cls.available_aliases()}")
+        return cls.REGISTRY[alias]
+
+    @classmethod
+    def get_model_class(cls, alias: Union[str, type]) -> type:
+        return cls._get_model_record(alias).model_class
+
+    @classmethod
+    def get_model_priority(cls, alias: Union[str, type]) -> int:
+        return cls._get_model_record(alias).ag_priority
+
+    @classmethod
+    def available_aliases(cls) -> list[str]:
+        return sorted(cls.REGISTRY.keys())

autogluon/timeseries/models/toto/__init__.py

@@ -0,0 +1,3 @@
+from .model import TotoModel
+
+__all__ = ["TotoModel"]

autogluon/timeseries/models/toto/_internal/__init__.py

@@ -0,0 +1,9 @@
+from .backbone import TotoBackbone
+from .dataset import MaskedTimeseries
+from .forecaster import TotoForecaster
+
+__all__ = [
+    "MaskedTimeseries",
+    "TotoBackbone",
+    "TotoForecaster",
+]

autogluon/timeseries/models/toto/_internal/backbone/__init__.py

@@ -0,0 +1,3 @@
+from .backbone import TotoBackbone
+
+__all__ = ["TotoBackbone"]