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

autogluon/timeseries/models/ensemble/abstract.py

@@ -1,9 +1,7 @@
-import functools
 import logging
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from typing import Optional
 
-import numpy as np
 from typing_extensions import final
 
 from autogluon.core.utils.exceptions import TimeLimitExceeded
@@ -18,32 +16,32 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
 
     @property
     @abstractmethod
-    def model_names(self) -> List[str]:
+    def model_names(self) -> list[str]:
         """Names of base models included in the ensemble."""
         pass
 
     @final
     def fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
     ):
         """Fit ensemble model given predictions of candidate base models and the true data.
 
         Parameters
         ----------
-        predictions_per_window : Dict[str, List[TimeSeriesDataFrame]]
+        predictions_per_window
             Dictionary that maps the names of component models to their respective predictions for each validation
             window.
-        data_per_window : List[TimeSeriesDataFrame]
+        data_per_window
             Observed ground truth data used to train the ensemble for each validation window. Each entry in the list
             includes both the forecast horizon (for which the predictions are given in ``predictions``), as well as the
             "history".
-        model_scores : Optional[Dict[str, float]]
+        model_scores
             Scores (higher is better) for the models that will constitute the ensemble.
-        time_limit : Optional[float]
+        time_limit
             Maximum allowed time for training in seconds.
         """
         if time_limit is not None and time_limit <= 0:
@@ -67,18 +65,18 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
 
     def _fit(
         self,
-        predictions_per_window: Dict[str, List[TimeSeriesDataFrame]],
-        data_per_window: List[TimeSeriesDataFrame],
-        model_scores: Optional[Dict[str, float]] = None,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
         time_limit: Optional[float] = None,
-    ):
+    ) -> None:
         """Private method for `fit`. See `fit` for documentation of arguments. Apart from the model
         training logic, `fit` additionally implements other logic such as keeping track of the time limit.
         """
         raise NotImplementedError
 
     @final
-    def predict(self, data: Dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+    def predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
         if not set(self.model_names).issubset(set(data.keys())):
             raise ValueError(
                 f"Set of models given for prediction in {self.name} differ from those provided during initialization."
@@ -93,47 +91,13 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         return self._predict(data=data, **kwargs)
 
     @abstractmethod
-    def _predict(self, data: Dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+    def _predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
         pass
 
     @abstractmethod
-    def remap_base_models(self, model_refit_map: Dict[str, str]) -> None:
+    def remap_base_models(self, model_refit_map: dict[str, str]) -> None:
         """Update names of the base models based on the mapping in model_refit_map.
 
         This method should be called after performing refit_full to point to the refitted base models, if necessary.
         """
         pass
-
-
-class AbstractWeightedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
-    """Abstract class for weighted ensembles which assign one (global) weight per model."""
-
-    def __init__(self, name: Optional[str] = None, **kwargs):
-        if name is None:
-            name = "WeightedEnsemble"
-        super().__init__(name=name, **kwargs)
-        self.model_to_weight: Dict[str, float] = {}
-
-    @property
-    def model_names(self) -> List[str]:
-        return list(self.model_to_weight.keys())
-
-    @property
-    def model_weights(self) -> np.ndarray:
-        return np.array(list(self.model_to_weight.values()), dtype=np.float64)
-
-    def _predict(self, data: Dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
-        weighted_predictions = [data[model_name] * weight for model_name, weight in self.model_to_weight.items()]
-        return functools.reduce(lambda x, y: x + y, weighted_predictions)
-
-    def get_info(self) -> dict:
-        info = super().get_info()
-        info["model_weights"] = self.model_to_weight.copy()
-        return info
-
-    def remap_base_models(self, model_refit_map: Dict[str, str]) -> None:
-        updated_weights = {}
-        for model, weight in self.model_to_weight.items():
-            model_full_name = model_refit_map.get(model, model)
-            updated_weights[model_full_name] = weight
-        self.model_to_weight = updated_weights

autogluon/timeseries/models/ensemble/array_based/__init__.py

@@ -0,0 +1,3 @@
+from .models import MedianEnsemble, PerQuantileTabularEnsemble, TabularEnsemble
+
+__all__ = ["MedianEnsemble", "PerQuantileTabularEnsemble", "TabularEnsemble"]

autogluon/timeseries/models/ensemble/array_based/abstract.py

@@ -0,0 +1,247 @@
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Sequence, Union
+
+import numpy as np
+from typing_extensions import Self
+
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+from autogluon.timeseries.metrics.abstract import TimeSeriesScorer
+from autogluon.timeseries.utils.features import CovariateMetadata
+
+from ..abstract import AbstractTimeSeriesEnsembleModel
+from .regressor import EnsembleRegressor
+
+
+class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
+    """Abstract base class for time series ensemble models which operate on arrays of base model
+    predictions for training and inference.
+
+    Other Parameters
+    ----------------
+    isotonization: str, default = "sort"
+        The isotonization method to use (i.e. the algorithm to prevent quantile non-crossing).
+        Currently only "sort" is supported.
+    detect_and_ignore_failures: bool, default = True
+        Whether to detect and ignore "failed models", defined as models which have a loss that is larger
+        than 10x the median loss of all the models. This can be very important for the regression-based
+        ensembles, as moving the weight from such a "failed model" to zero can require a long training
+        time.
+    """
+
+    def __init__(
+        self,
+        path: Optional[str] = None,
+        name: Optional[str] = None,
+        hyperparameters: Optional[dict[str, Any]] = None,
+        freq: Optional[str] = None,
+        prediction_length: int = 1,
+        covariate_metadata: Optional[CovariateMetadata] = 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,
+    ):
+        super().__init__(
+            path=path,
+            name=name,
+            hyperparameters=hyperparameters,
+            freq=freq,
+            prediction_length=prediction_length,
+            covariate_metadata=covariate_metadata,
+            target=target,
+            quantile_levels=quantile_levels,
+            eval_metric=eval_metric,
+        )
+        self.ensemble_regressor: Optional[EnsembleRegressor] = None
+        self._model_names: list[str] = []
+
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
+        return {
+            "isotonization": "sort",
+            "detect_and_ignore_failures": True,
+        }
+
+    @classmethod
+    def load(cls, path: str, reset_paths: bool = True, load_oof: bool = False, verbose: bool = True) -> Self:
+        model = super().load(path=path, reset_paths=reset_paths, load_oof=load_oof, verbose=verbose)
+
+        if reset_paths and model.ensemble_regressor is not None:
+            model.ensemble_regressor.set_path(os.path.join(model.path, "ensemble_regressor"))
+
+        return model
+
+    @staticmethod
+    def to_array(df: TimeSeriesDataFrame) -> np.ndarray:
+        """Given a TimeSeriesDataFrame object, return a single array composing the values contained
+        in the data frame.
+
+        Parameters
+        ----------
+        df
+            TimeSeriesDataFrame to convert to an array. Must contain exactly `prediction_length`
+            values for each item. The columns of `df` can correspond to ground truth values
+            or predictions (in which case, these will be the mean or quantile forecasts).
+
+        Returns
+        -------
+        array
+            of shape (num_items, prediction_length, num_outputs).
+        """
+        assert df.index.is_monotonic_increasing
+        array = df.to_numpy()
+        num_items = df.num_items
+        shape = (
+            num_items,
+            df.shape[0] // num_items,  # timesteps per item
+            df.shape[1],  # num_outputs
+        )
+        return array.reshape(shape)
+
+    def _get_base_model_predictions(
+        self,
+        predictions_per_window: Union[dict[str, list[TimeSeriesDataFrame]], dict[str, TimeSeriesDataFrame]],
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Given a mapping from model names to a list of data frames representing
+        their predictions per window, return a multidimensional array representation.
+
+        Parameters
+        ----------
+        predictions_per_window
+            A dictionary with list[TimeSeriesDataFrame] values, where each TimeSeriesDataFrame
+            contains predictions for the window in question. If the dictionary values are
+            TimeSeriesDataFrame, they will be treated like a single window.
+
+        Returns
+        -------
+        base_model_mean_predictions
+            Array of shape (num_windows, num_items, prediction_length, 1, num_models)
+        base_model_quantile_predictions
+            Array of shape (num_windows, num_items, prediction_length, num_quantiles, num_models)
+        """
+
+        if not predictions_per_window:
+            raise ValueError("No base model predictions are provided.")
+
+        first_prediction = list(predictions_per_window.values())[0]
+        if isinstance(first_prediction, TimeSeriesDataFrame):
+            predictions_per_window = {k: [v] for k, v in predictions_per_window.items()}  # type: ignore
+
+        predictions = {
+            model_name: [self.to_array(window) for window in windows]  # type: ignore
+            for model_name, windows in predictions_per_window.items()
+        }
+        base_model_predictions = np.stack([x for x in predictions.values()], axis=-1)
+
+        return base_model_predictions[:, :, :, :1, :], base_model_predictions[:, :, :, 1:, :]
+
+    def _isotonize(self, prediction_array: np.ndarray) -> np.ndarray:
+        """Apply isotonization to ensure quantile non-crossing.
+
+        Parameters
+        ----------
+        prediction_array
+            Array of shape (num_windows, num_items, prediction_length, num_quantiles)
+
+        Returns
+        -------
+        isotonized_array
+            Array with same shape but quantiles sorted along last dimension
+        """
+        isotonization = self.get_hyperparameters()["isotonization"]
+        if isotonization == "sort":
+            return np.sort(prediction_array, axis=-1)
+        return prediction_array
+
+    def _fit(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: Optional[dict[str, float]] = None,
+        time_limit: Optional[float] = None,
+    ) -> None:
+        # process inputs
+        filtered_predictions = self._filter_failed_models(predictions_per_window, model_scores)
+        base_model_mean_predictions, base_model_quantile_predictions = self._get_base_model_predictions(
+            filtered_predictions
+        )
+
+        # process labels
+        ground_truth_per_window = [y.slice_by_timestep(-self.prediction_length, None) for y in data_per_window]
+        labels = np.stack(
+            [self.to_array(gt) for gt in ground_truth_per_window], axis=0
+        )  # (num_windows, num_items, prediction_length, 1)
+
+        self._model_names = list(filtered_predictions.keys())
+        self.ensemble_regressor = self._get_ensemble_regressor()
+        self.ensemble_regressor.fit(
+            base_model_mean_predictions=base_model_mean_predictions,
+            base_model_quantile_predictions=base_model_quantile_predictions,
+            labels=labels,
+            time_limit=time_limit,
+        )
+
+    @abstractmethod
+    def _get_ensemble_regressor(self) -> EnsembleRegressor:
+        pass
+
+    def _predict(self, data: dict[str, TimeSeriesDataFrame], **kwargs) -> TimeSeriesDataFrame:
+        if self.ensemble_regressor is None:
+            if not self._model_names:
+                raise ValueError("Ensemble model has not been fitted yet.")
+            # Try to recreate the regressor (for loaded models)
+            self.ensemble_regressor = self._get_ensemble_regressor()
+
+        input_data = {}
+        for m in self.model_names:
+            assert m in data, f"Predictions for model {m} not provided during ensemble prediction."
+            input_data[m] = data[m]
+
+        base_model_mean_predictions, base_model_quantile_predictions = self._get_base_model_predictions(input_data)
+
+        mean_predictions, quantile_predictions = self.ensemble_regressor.predict(
+            base_model_mean_predictions=base_model_mean_predictions,
+            base_model_quantile_predictions=base_model_quantile_predictions,
+        )
+
+        quantile_predictions = self._isotonize(quantile_predictions)
+        prediction_array = np.concatenate([mean_predictions, quantile_predictions], axis=-1)
+
+        output = list(input_data.values())[0].copy()
+        num_folds, num_items, num_timesteps, num_outputs = prediction_array.shape
+        assert (num_folds, num_timesteps) == (1, self.prediction_length)
+        assert len(output.columns) == num_outputs
+
+        output[output.columns] = prediction_array.reshape((num_items * num_timesteps, num_outputs))
+
+        return output
+
+    @property
+    def model_names(self) -> list[str]:
+        return self._model_names
+
+    def remap_base_models(self, model_refit_map: dict[str, str]) -> None:
+        """Update names of the base models based on the mapping in model_refit_map."""
+        self._model_names = [model_refit_map.get(name, name) for name in self._model_names]
+
+    def _filter_failed_models(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        model_scores: Optional[dict[str, float]],
+    ) -> dict[str, list[TimeSeriesDataFrame]]:
+        """Filter out failed models based on detect_and_ignore_failures setting."""
+        if not self.get_hyperparameters()["detect_and_ignore_failures"]:
+            return predictions_per_window
+
+        if model_scores is None or len(model_scores) == 0:
+            return predictions_per_window
+
+        valid_scores = {k: v for k, v in model_scores.items() if np.isfinite(v)}
+        if len(valid_scores) == 0:
+            raise ValueError("All models have NaN scores. At least one model must run successfully to fit an ensemble")
+
+        losses = {k: -v for k, v in valid_scores.items()}
+        median_loss = np.nanmedian(list(losses.values()))
+        threshold = 10 * median_loss
+        good_models = {k for k, loss in losses.items() if loss <= threshold}
+
+        return {k: v for k, v in predictions_per_window.items() if k in good_models}

autogluon/timeseries/models/ensemble/array_based/models.py

@@ -0,0 +1,50 @@
+import os
+from abc import ABC
+from typing import Any, Type
+
+from .abstract import ArrayBasedTimeSeriesEnsembleModel
+from .regressor import (
+    EnsembleRegressor,
+    MedianEnsembleRegressor,
+    PerQuantileTabularEnsembleRegressor,
+    TabularEnsembleRegressor,
+)
+
+
+class MedianEnsemble(ArrayBasedTimeSeriesEnsembleModel):
+    def _get_ensemble_regressor(self) -> MedianEnsembleRegressor:
+        return MedianEnsembleRegressor()
+
+
+class BaseTabularEnsemble(ArrayBasedTimeSeriesEnsembleModel, ABC):
+    ensemble_regressor_type: Type[EnsembleRegressor]
+
+    def _get_default_hyperparameters(self) -> dict[str, Any]:
+        default_hps = super()._get_default_hyperparameters()
+        default_hps.update(
+            {
+                "tabular_hyperparameters": {"GBM": {}},
+            }
+        )
+        return default_hps
+
+    def _get_ensemble_regressor(self):
+        return self.ensemble_regressor_type(
+            path=os.path.join(self.path, "ensemble_regressor"),
+            quantile_levels=list(self.quantile_levels),
+            tabular_hyperparameters=self.get_hyperparameters()["tabular_hyperparameters"],
+        )
+
+
+class TabularEnsemble(BaseTabularEnsemble):
+    """Time series ensemble model using single AutoGluon TabularPredictor for all quantiles."""
+
+    ensemble_regressor_type = TabularEnsembleRegressor
+
+
+class PerQuantileTabularEnsemble(BaseTabularEnsemble):
+    """Time series ensemble model using separate `TabularPredictor` instances for each quantile in
+    addition to a dedicated `TabularPredictor` for the mean (point) forecast.
+    """
+
+    ensemble_regressor_type = PerQuantileTabularEnsembleRegressor

autogluon/timeseries/models/ensemble/array_based/regressor/__init__.py

@@ -0,0 +1,10 @@
+from .abstract import EnsembleRegressor, MedianEnsembleRegressor
+from .per_quantile_tabular import PerQuantileTabularEnsembleRegressor
+from .tabular import TabularEnsembleRegressor
+
+__all__ = [
+    "EnsembleRegressor",
+    "MedianEnsembleRegressor",
+    "PerQuantileTabularEnsembleRegressor",
+    "TabularEnsembleRegressor",
+]

autogluon/timeseries/models/ensemble/array_based/regressor/abstract.py

@@ -0,0 +1,87 @@
+from abc import ABC, abstractmethod
+
+import numpy as np
+from typing_extensions import Self
+
+
+class EnsembleRegressor(ABC):
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def set_path(self, path: str) -> None:
+        pass
+
+    @abstractmethod
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        **kwargs,
+    ) -> Self:
+        """
+        Parameters
+        ----------
+        base_model_mean_predictions
+            Mean (point) predictions of base models. Array of shape
+            (num_windows, num_items, prediction_length, 1, num_models)
+
+        base_model_quantile_predictions
+            Quantile predictions of base models. Array of shape
+            (num_windows, num_items, prediction_length, num_quantiles, num_models)
+
+        labels
+            Ground truth array of shape
+            (num_windows, num_items, prediction_length, 1)
+        """
+        pass
+
+    @abstractmethod
+    def predict(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Predict with the fitted ensemble regressor for a single window.
+        The items do not have to refer to the same item indices used when fitting
+        the model.
+
+        Parameters
+        ----------
+        base_model_mean_predictions
+            Mean (point) predictions of base models. Array of shape
+            (1, num_items, prediction_length, 1, num_models)
+
+        base_model_quantile_predictions
+            Quantile predictions of base models. Array of shape
+            (1, num_items, prediction_length, num_quantiles, num_models)
+
+        Returns
+        -------
+        ensemble_mean_predictions
+            Array of shape (1, num_items, prediction_length, 1)
+        ensemble_quantile_predictions
+            Array of shape (1, num_items, prediction_length, num_quantiles)
+        """
+        pass
+
+
+class MedianEnsembleRegressor(EnsembleRegressor):
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        **kwargs,
+    ) -> Self:
+        return self
+
+    def predict(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        return (
+            np.nanmedian(base_model_mean_predictions, axis=-1),
+            np.nanmedian(base_model_quantile_predictions, axis=-1),
+        )

autogluon/timeseries/models/ensemble/array_based/regressor/per_quantile_tabular.py

@@ -0,0 +1,133 @@
+import logging
+import os
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+from typing_extensions import Self
+
+from autogluon.tabular import TabularPredictor
+
+from .abstract import EnsembleRegressor
+
+logger = logging.getLogger(__name__)
+
+
+class PerQuantileTabularEnsembleRegressor(EnsembleRegressor):
+    """TabularPredictor ensemble regressor using separate models per quantile plus dedicated mean model."""
+
+    def __init__(
+        self,
+        path: str,
+        quantile_levels: list[float],
+        tabular_hyperparameters: Optional[dict] = None,
+    ):
+        super().__init__()
+        self.path = path
+        self.quantile_levels = quantile_levels
+        self.tabular_hyperparameters = tabular_hyperparameters or {}
+        self.quantile_predictors: list[TabularPredictor] = []
+        self.mean_predictor: Optional[TabularPredictor] = None
+
+    def set_path(self, path: str) -> None:
+        self.path = path
+
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        **kwargs,
+    ) -> Self:
+        """Fit separate TabularPredictor for mean and each quantile level."""
+        # TODO: implement time_limit
+
+        num_windows, num_items, prediction_length = base_model_mean_predictions.shape[:3]
+        target = labels.reshape(num_windows * num_items * prediction_length).ravel()
+
+        # fit mean predictor, based on mean predictions of base models
+        mean_df = self._get_feature_df(base_model_mean_predictions, 0)
+        mean_df["target"] = target
+        self.mean_predictor = TabularPredictor(
+            label="target",
+            path=os.path.join(self.path, "mean"),
+            verbosity=1,
+            problem_type="regression",
+        ).fit(
+            mean_df,
+            hyperparameters=self.tabular_hyperparameters,
+        )
+
+        # fit quantile predictors, each quantile predictor is based on the
+        # estimates of that quantile from base models
+        for i, quantile in enumerate(self.quantile_levels):
+            q_df = self._get_feature_df(base_model_quantile_predictions, i)
+            q_df["target"] = target
+
+            predictor = TabularPredictor(
+                label="target",
+                path=os.path.join(self.path, f"quantile_{quantile}"),
+                verbosity=1,
+                problem_type="regression",
+            ).fit(q_df, hyperparameters=self.tabular_hyperparameters)
+            self.quantile_predictors.append(predictor)
+
+        return self
+
+    def _get_feature_df(self, predictions: np.ndarray, index: int) -> pd.DataFrame:
+        num_windows, num_items, prediction_length, _, num_models = predictions.shape
+        num_tabular_items = num_windows * num_items * prediction_length
+
+        df = pd.DataFrame(
+            predictions[:, :, :, index].reshape(num_tabular_items, num_models),
+            columns=[f"model_{mi}" for mi in range(num_models)],
+        )
+
+        return df
+
+    def load_predictors(self):
+        if self.mean_predictor is None or len(self.quantile_predictors) < len(self.quantile_levels):
+            try:
+                self.mean_predictor = TabularPredictor.load(os.path.join(self.path, "mean"))
+
+                self.quantile_predictors = []
+                for quantile in self.quantile_levels:
+                    predictor = TabularPredictor.load(os.path.join(self.path, f"quantile_{quantile}"))
+                    self.quantile_predictors.append(predictor)
+
+            except FileNotFoundError:
+                raise ValueError("Model must be fitted before loading for prediction")
+
+    def predict(
+        self, base_model_mean_predictions: np.ndarray, base_model_quantile_predictions: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray]:
+        self.load_predictors()
+
+        num_windows, num_items, prediction_length, _, _ = base_model_mean_predictions.shape
+        assert num_windows == 1, "Prediction expects a single window to be provided"
+
+        # predict means
+        assert self.mean_predictor is not None
+        mean_predictions = self.mean_predictor.predict(
+            self._get_feature_df(base_model_mean_predictions, 0),
+            as_pandas=False,
+        ).reshape(num_windows, num_items, prediction_length, 1)
+
+        # predict quantiles
+        quantile_predictions_list = []
+        for i, predictor in enumerate(self.quantile_predictors):
+            quantile_predictions_list.append(
+                predictor.predict(self._get_feature_df(base_model_quantile_predictions, i), as_pandas=False).reshape(
+                    num_windows, num_items, prediction_length
+                )
+            )
+        quantile_predictions = np.stack(quantile_predictions_list, axis=-1)
+
+        return mean_predictions, quantile_predictions
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        # Remove predictors to avoid pickling heavy TabularPredictor objects
+        state["mean_predictor"] = None
+        state["quantile_predictors"] = []
+        return state