autogluon.timeseries 1.4.1b20251115__py3-none-any.whl → 1.5.0b20251221__py3-none-any.whl

autogluon/timeseries/models/ensemble/abstract.py

@@ -1,6 +1,5 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import Optional
 
 from typing_extensions import final
 
@@ -12,7 +11,12 @@ logger = logging.getLogger(__name__)
 
 
 class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
-    """Abstract class for time series ensemble models."""
+    """Abstract base class for time series ensemble models that combine predictions from multiple base models.
+
+    Ensemble training process operates on validation predictions from base models rather than raw time series
+    data. This allows the ensemble to learn optimal combination strategies based on each model's performance
+    across different validation windows and time series patterns.
+    """
 
     @property
     @abstractmethod
@@ -25,8 +29,8 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         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,
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
     ):
         """Fit ensemble model given predictions of candidate base models and the true data.
 
@@ -50,7 +54,7 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
             )
             raise TimeLimitExceeded
         if isinstance(data_per_window, TimeSeriesDataFrame):
-            raise ValueError("When fitting ensemble, `data` should contain ground truth for each validation window")
+            raise ValueError("When fitting ensemble, ``data`` should contain ground truth for each validation window")
         num_val_windows = len(data_per_window)
         for model, preds in predictions_per_window.items():
             if len(preds) != num_val_windows:
@@ -67,11 +71,11 @@ class AbstractTimeSeriesEnsembleModel(TimeSeriesModelBase, ABC):
         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,
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = 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.
+        """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
 

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

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

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

@@ -1,9 +1,7 @@
-import os
 from abc import ABC, abstractmethod
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Sequence
 
 import numpy as np
-from typing_extensions import Self
 
 from autogluon.timeseries.dataset import TimeSeriesDataFrame
 from autogluon.timeseries.metrics.abstract import TimeSeriesScorer
@@ -14,15 +12,19 @@ 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.
+    """Abstract base class for ensemble models that operate on multi-dimensional arrays of base model predictions.
+
+    Array-based ensembles convert time series predictions into structured numpy arrays for efficient processing
+    and enable sophisticated combination strategies beyond simple weighted averaging. Array-based ensembles also
+    support isotonization in quantile forecasts--ensuring quantile crossing does not occur. They also have built-in
+    failed model detection and filtering capabilities.
 
     Other Parameters
     ----------------
-    isotonization: str, default = "sort"
+    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
+    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
@@ -31,15 +33,15 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, 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,
     ):
         super().__init__(
             path=path,
@@ -52,7 +54,7 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
             quantile_levels=quantile_levels,
             eval_metric=eval_metric,
         )
-        self.ensemble_regressor: Optional[EnsembleRegressor] = None
+        self.ensemble_regressor: EnsembleRegressor | None = None
         self._model_names: list[str] = []
 
     def _get_default_hyperparameters(self) -> dict[str, Any]:
@@ -61,15 +63,6 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
             "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
@@ -78,8 +71,8 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
         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
+            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
@@ -99,7 +92,7 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
 
     def _get_base_model_predictions(
         self,
-        predictions_per_window: Union[dict[str, list[TimeSeriesDataFrame]], dict[str, TimeSeriesDataFrame]],
+        predictions_per_window: 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.
@@ -147,7 +140,7 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
         isotonized_array
             Array with same shape but quantiles sorted along last dimension
         """
-        isotonization = self.get_hyperparameters()["isotonization"]
+        isotonization = self.get_hyperparameter("isotonization")
         if isotonization == "sort":
             return np.sort(prediction_array, axis=-1)
         return prediction_array
@@ -156,8 +149,8 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
         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,
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
     ) -> None:
         # process inputs
         filtered_predictions = self._filter_failed_models(predictions_per_window, model_scores)
@@ -226,10 +219,10 @@ class ArrayBasedTimeSeriesEnsembleModel(AbstractTimeSeriesEnsembleModel, ABC):
     def _filter_failed_models(
         self,
         predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
-        model_scores: Optional[dict[str, float]],
+        model_scores: dict[str, float] | None,
     ) -> dict[str, list[TimeSeriesDataFrame]]:
         """Filter out failed models based on detect_and_ignore_failures setting."""
-        if not self.get_hyperparameters()["detect_and_ignore_failures"]:
+        if not self.get_hyperparameter("detect_and_ignore_failures"):
             return predictions_per_window
 
         if model_scores is None or len(model_scores) == 0:

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

@@ -1,10 +1,12 @@
-import os
 from abc import ABC
 from typing import Any, Type
 
+from autogluon.timeseries.dataset import TimeSeriesDataFrame
+
 from .abstract import ArrayBasedTimeSeriesEnsembleModel
 from .regressor import (
     EnsembleRegressor,
+    LinearStackerEnsembleRegressor,
     MedianEnsembleRegressor,
     PerQuantileTabularEnsembleRegressor,
     TabularEnsembleRegressor,
@@ -12,6 +14,21 @@ from .regressor import (
 
 
 class MedianEnsemble(ArrayBasedTimeSeriesEnsembleModel):
+    """Robust ensemble that computes predictions as the element-wise median of base model mean
+    and quantile forecasts, providing robustness to outlier predictions.
+
+    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 _get_ensemble_regressor(self) -> MedianEnsembleRegressor:
         return MedianEnsembleRegressor()
 
@@ -21,30 +38,148 @@ class BaseTabularEnsemble(ArrayBasedTimeSeriesEnsembleModel, ABC):
 
     def _get_default_hyperparameters(self) -> dict[str, Any]:
         default_hps = super()._get_default_hyperparameters()
-        default_hps.update(
-            {
-                "tabular_hyperparameters": {"GBM": {}},
-            }
-        )
+        default_hps.update({"model_name": "CAT", "model_hyperparameters": {}})
         return default_hps
 
     def _get_ensemble_regressor(self):
+        hyperparameters = self.get_hyperparameters()
         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"],
+            model_name=hyperparameters["model_name"],
+            model_hyperparameters=hyperparameters["model_hyperparameters"],
         )
 
 
 class TabularEnsemble(BaseTabularEnsemble):
-    """Time series ensemble model using single AutoGluon TabularPredictor for all quantiles."""
+    """Tabular ensemble that uses a single AutoGluon-Tabular model to learn ensemble combinations.
+
+    This ensemble trains a single tabular model (such as gradient boosting machines) to predict all
+    quantiles simultaneously from base model predictions. The tabular model learns complex non-linear
+    patterns in how base models should be combined, potentially capturing interactions and conditional
+    dependencies that simple weighted averages cannot represent.
+
+    Other Parameters
+    ----------------
+    model_name : str, default = "CAT"
+        Name of the AutoGluon-Tabular model to use for ensemble learning. Model name should be registered
+        in AutoGluon-Tabular model registry.
+    model_hyperparameters : dict, default = {}
+        Hyperparameters to pass to the underlying AutoGluon-Tabular model.
+    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.
+    """
 
     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.
+    """Tabular ensemble using separate AutoGluon-Tabular models for each quantile and mean forecast.
+
+    This ensemble trains dedicated tabular models for each quantile level plus a separate model
+    for the mean prediction. Each model specializes in learning optimal combinations for its
+    specific target, allowing for quantile-specific ensemble strategies that can capture different
+    model behaviors across the prediction distribution.
+
+    Other Parameters
+    ----------------
+    model_name : str, default = "GBM"
+        Name of the AutoGluon-Tabular model to use for ensemble learning. Model name should be registered
+        in AutoGluon-Tabular model registry.
+    model_hyperparameters : dict, default = {}
+        Hyperparameters to pass to the underlying AutoGluon-Tabular model.
+    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.
     """
 
     ensemble_regressor_type = PerQuantileTabularEnsembleRegressor
+
+
+class LinearStackerEnsemble(ArrayBasedTimeSeriesEnsembleModel):
+    """Linear stacking ensemble that learns optimal linear combination weights through gradient-based
+    optimization.
+
+    Weighted combinations can be per model or per model-quantile, model-horizon, model-quantile-horizon
+    combinations. These choices are controlled by the ``weights_per`` hyperparameter.
+
+    The optimization process uses gradient descent with configurable learning rates and convergence
+    criteria, allowing for flexible training dynamics. Weight pruning can be applied to remove
+    models with negligible contributions, resulting in sparse and interpretable ensembles.
+
+    Other Parameters
+    ----------------
+    weights_per : str, default = "m"
+        Granularity of weight learning.
+
+        - "m": single weight per model
+        - "mq": single weight for each model-quantile combination
+        - "mt": single weight for each model-time step where time steps run across the prediction horizon
+        - "mtq": single weight for each model-quantile-time step combination
+    lr : float, default = 0.1
+        Learning rate for PyTorch optimizer during weight training.
+    max_epochs : int, default = 10000
+        Maximum number of training epochs for weight optimization.
+    relative_tolerance : float, default = 1e-7
+        Relative tolerance for convergence detection during training.
+    prune_below : float, default = 0.0
+        Threshold below which weights are pruned to zero for sparsity. The weights are redistributed across
+        remaining models after pruning.
+    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 _get_default_hyperparameters(self) -> dict[str, Any]:
+        default_hps = super()._get_default_hyperparameters()
+        default_hps.update(
+            {
+                "weights_per": "m",
+                "lr": 0.1,
+                "max_epochs": 10000,
+                "relative_tolerance": 1e-7,
+                "prune_below": 0.0,
+            }
+        )
+        return default_hps
+
+    def _get_ensemble_regressor(self) -> LinearStackerEnsembleRegressor:
+        hps = self.get_hyperparameters()
+        return LinearStackerEnsembleRegressor(
+            quantile_levels=list(self.quantile_levels),
+            weights_per=hps["weights_per"],
+            lr=hps["lr"],
+            max_epochs=hps["max_epochs"],
+            relative_tolerance=hps["relative_tolerance"],
+            prune_below=hps["prune_below"],
+        )
+
+    def _fit(
+        self,
+        predictions_per_window: dict[str, list[TimeSeriesDataFrame]],
+        data_per_window: list[TimeSeriesDataFrame],
+        model_scores: dict[str, float] | None = None,
+        time_limit: float | None = None,
+    ) -> None:
+        super()._fit(predictions_per_window, data_per_window, model_scores, time_limit)
+
+        assert isinstance(self.ensemble_regressor, LinearStackerEnsembleRegressor)
+
+        if self.ensemble_regressor.kept_indices is not None:
+            original_names = self._model_names
+            self._model_names = [original_names[i] for i in self.ensemble_regressor.kept_indices]

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

@@ -1,9 +1,11 @@
 from .abstract import EnsembleRegressor, MedianEnsembleRegressor
+from .linear_stacker import LinearStackerEnsembleRegressor
 from .per_quantile_tabular import PerQuantileTabularEnsembleRegressor
 from .tabular import TabularEnsembleRegressor
 
 __all__ = [
     "EnsembleRegressor",
+    "LinearStackerEnsembleRegressor",
     "MedianEnsembleRegressor",
     "PerQuantileTabularEnsembleRegressor",
     "TabularEnsembleRegressor",

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

@@ -8,16 +8,13 @@ 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,
+        time_limit: float | None = None,
     ) -> Self:
         """
         Parameters
@@ -33,6 +30,10 @@ class EnsembleRegressor(ABC):
         labels
             Ground truth array of shape
             (num_windows, num_items, prediction_length, 1)
+
+        time_limit
+            Approximately how long ``fit`` will run (wall-clock time in seconds). If
+            not specified, training time will not be limited.
         """
         pass
 
@@ -72,7 +73,7 @@ class MedianEnsembleRegressor(EnsembleRegressor):
         base_model_mean_predictions: np.ndarray,
         base_model_quantile_predictions: np.ndarray,
         labels: np.ndarray,
-        **kwargs,
+        time_limit: float | None = None,
     ) -> Self:
         return self
 

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

@@ -0,0 +1,186 @@
+from typing import Literal
+
+import numpy as np
+from typing_extensions import Self
+
+from autogluon.timeseries.utils.timer import Timer
+
+from .abstract import EnsembleRegressor
+
+
+class LinearStackerEnsembleRegressor(EnsembleRegressor):
+    """Linear stacker ensemble regressor using PyTorch optimization with softmax weights.
+
+    Implements weighted averaging of base model predictions with learnable weights optimized
+    via gradient descent. Uses PyTorch during training for optimization, then stores weights
+    as numpy arrays for efficient prediction.
+
+    Parameters
+    ----------
+    quantile_levels
+        List of quantile levels for quantile predictions (e.g., [0.1, 0.5, 0.9]).
+    weights_per
+        Weight configuration specifying which dimensions to learn weights for:
+
+        - "m": Per-model weights (shape: num_models), defaults to "m"
+        - "mt": Per-model and per-time weights (shape: prediction_length, num_models)
+        - "mq": Per-model and per-model-output (quantiles and mean) weights
+          (shape: num_quantiles+1, num_models)
+        - "mtq": Per-model, per-time, and per-quantile weights
+          (shape: prediction_length, num_quantiles+1, num_models)
+    lr
+        Learning rate for Adam optimizer. Defaults to 0.1.
+    max_epochs
+        Maximum number of training epochs. Defaults to 10000.
+    relative_tolerance
+        Convergence tolerance for relative loss change between epochs. Defaults to 1e-7.
+    prune_below
+        Importance threshold for model sparsification. Models with importance below this
+        threshold are dropped after weight optimization. Set to 0.0 to disable sparsification.
+        Defaults to 0.0.
+    """
+
+    def __init__(
+        self,
+        quantile_levels: list[float],
+        weights_per: Literal["m", "mt", "mq", "mtq"] = "m",
+        lr: float = 0.1,
+        max_epochs: int = 10_000,
+        relative_tolerance: float = 1e-7,
+        prune_below: float = 0.0,
+    ):
+        super().__init__()
+        self.quantile_levels = quantile_levels
+        self.weights_per = weights_per
+        self.lr = lr
+        self.max_epochs = max_epochs
+        self.relative_tolerance = relative_tolerance
+        self.prune_below = prune_below
+
+        self.weights: np.ndarray | None = None
+        self.kept_indices: list[int] | None = None
+
+    def _compute_weight_shape(self, base_model_predictions_shape: tuple) -> tuple:
+        """Compute weight tensor shape based on weights_per configuration."""
+        _, _, prediction_length, num_outputs, num_models = base_model_predictions_shape
+
+        shapes = {
+            "m": (1, 1, num_models),
+            "mt": (prediction_length, 1, num_models),
+            "mq": (1, num_outputs, num_models),
+            "mtq": (prediction_length, num_outputs, num_models),
+        }
+        try:
+            return (1, 1) + shapes[self.weights_per]
+        except KeyError:
+            raise ValueError(f"Unsupported weights_per: {self.weights_per}")
+
+    def make_weighted_average_module(self, base_model_predictions_shape: tuple):
+        import torch
+
+        class WeightedAverage(torch.nn.Module):
+            def __init__(self, shape):
+                super().__init__()
+                self.raw_weights = torch.nn.Parameter(torch.zeros(*shape, dtype=torch.float32))
+
+            def get_normalized_weights(self):
+                return torch.softmax(self.raw_weights, dim=-1)  # softmax over models
+
+            def forward(self, base_model_predictions: torch.Tensor):
+                return torch.sum(self.get_normalized_weights() * base_model_predictions, dim=-1)
+
+        return WeightedAverage(self._compute_weight_shape(base_model_predictions_shape))
+
+    def fit(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+        labels: np.ndarray,
+        time_limit: float | None = None,
+    ) -> Self:
+        import torch
+
+        def _ql(
+            labels_tensor: torch.Tensor,
+            ensemble_predictions: torch.Tensor,
+        ) -> torch.Tensor:
+            """Compute the weighted quantile loss on predictions and ground truth (labels).
+            Considering that the first dimension of predictions is the mean, we treat
+            mean predictions on the same footing as median (0.5) predictions as contribution
+            to the overall weighted quantile loss.
+            """
+            quantile_levels = torch.tensor([0.5] + self.quantile_levels, dtype=torch.float32)
+            error = labels_tensor - ensemble_predictions  # (num_windows, num_items, num_time, num_outputs)
+            quantile_loss = torch.maximum(quantile_levels * error, (quantile_levels - 1) * error)
+            return torch.mean(quantile_loss)
+
+        timer = Timer(time_limit).start()
+
+        base_model_predictions = torch.tensor(
+            np.concatenate(
+                [base_model_mean_predictions, base_model_quantile_predictions],
+                axis=3,
+            ),
+            dtype=torch.float32,
+        )
+        labels_tensor = torch.tensor(labels, dtype=torch.float32)
+
+        weighted_average = self.make_weighted_average_module(base_model_predictions.shape)
+
+        optimizer = torch.optim.Adam(weighted_average.parameters(), lr=self.lr)
+
+        prev_loss = float("inf")
+        for _ in range(self.max_epochs):
+            optimizer.zero_grad()
+
+            ensemble_predictions = weighted_average(base_model_predictions)
+
+            loss = _ql(labels_tensor, ensemble_predictions)
+            loss.backward()
+            optimizer.step()
+
+            loss_change = abs(prev_loss - loss.item()) / (loss.item() + 1e-8)
+            if loss_change < self.relative_tolerance:
+                break
+            prev_loss = loss.item()
+
+            if timer.timed_out():
+                break
+
+        with torch.no_grad():
+            self.weights = weighted_average.get_normalized_weights().detach().numpy()
+
+        assert self.weights is not None
+        if self.prune_below > 0.0:
+            importances = self.weights.mean(axis=tuple(range(self.weights.ndim - 1)))  # shape (num_models,)
+
+            mask = importances >= self.prune_below
+            if not mask.any():
+                mask[importances.argmax()] = True
+
+            if not mask.all():
+                self.kept_indices = np.where(mask)[0].tolist()
+                self.weights = self.weights[..., mask]
+                self.weights = self.weights / self.weights.sum(axis=-1, keepdims=True)
+
+        return self
+
+    def predict(
+        self,
+        base_model_mean_predictions: np.ndarray,
+        base_model_quantile_predictions: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if self.weights is None:
+            raise ValueError("Model must be fitted before prediction")
+
+        all_predictions = np.concatenate([base_model_mean_predictions, base_model_quantile_predictions], axis=3)
+
+        if self.kept_indices is not None:
+            assert all_predictions.shape[-1] == len(self.kept_indices)
+
+        ensemble_pred = np.sum(self.weights * all_predictions, axis=-1)
+
+        mean_predictions = ensemble_pred[:, :, :, :1]
+        quantile_predictions = ensemble_pred[:, :, :, 1:]
+
+        return mean_predictions, quantile_predictions