remove quantile support from Timexer

Signed-off-by: Faakhir30 <zahidfaakhir@gmail.com>

Author: Faakhir30

pytorch_forecasting/layers/_output/_flatten_head.py

@@ -19,19 +19,13 @@ class FlattenHead(nn.Module):
         nf (int): Number of features in the last layer.
         target_window (int): Target window size.
         head_dropout (float): Dropout rate for the head. Defaults to 0.
-        n_quantiles (int, optional): Number of quantiles. Defaults to None."""
+    """
 
-    def __init__(self, n_vars, nf, target_window, head_dropout=0, n_quantiles=None):
+    def __init__(self, n_vars, nf, target_window, head_dropout=0):
         super().__init__()
         self.n_vars = n_vars
         self.flatten = nn.Flatten(start_dim=-2)
         self.linear = nn.Linear(nf, target_window)
-        self.n_quantiles = n_quantiles
-
-        if self.n_quantiles is not None:
-            self.linear = nn.Linear(nf, target_window * n_quantiles)
-        else:
-            self.linear = nn.Linear(nf, target_window)
         self.dropout = nn.Dropout(head_dropout)
 
     def forward(self, x):
@@ -40,7 +34,6 @@ def forward(self, x):
         x = self.dropout(x)
         x = x.permute(0, 2, 1)
 
-        if self.n_quantiles is not None:
-            batch_size = x.shape[0]
-            x = x.reshape(batch_size, -1, self.n_quantiles)
+        batch_size = x.shape[0]
+        x = x.reshape(batch_size, -1)
         return x

pytorch_forecasting/models/timexer/_timexer.py

@@ -18,11 +18,9 @@
 from pytorch_forecasting.metrics import (
     MAE,
     MAPE,
-    MASE,
     RMSE,
     SMAPE,
     MultiHorizonMetric,
-    QuantileLoss,
 )
 from pytorch_forecasting.metrics.base_metrics import MultiLoss
 from pytorch_forecasting.models.base import BaseModelWithCovariates
@@ -222,16 +220,6 @@ def __init__(
         if enc_in is None:
             self.enc_in = len(self.reals)
 
-        # NOTE: assume point prediction as default here,
-        # with single median quantile being the point prediction.
-        # hence self.n_quantiles = 1 for point predictions.
-        self.n_quantiles = 1
-
-        # set n_quantiles to the length of the quantiles list passed
-        # into the "quantiles" parameter when QuantileLoss is used.
-        if isinstance(loss, QuantileLoss):
-            self.n_quantiles = len(loss.quantiles)
-
         if hidden_size % n_heads != 0:
             raise ValueError(
                 f"hidden_size ({hidden_size}) must be divisible by n_heads ({n_heads}) "
@@ -302,7 +290,6 @@ def __init__(
             self.head_nf,
             self.hparams.prediction_length,
             head_dropout=self.hparams.dropout,
-            n_quantiles=self.n_quantiles,
         )
 
     @classmethod
@@ -507,11 +494,7 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
             if prediction.size(2) != len(target_positions):
                 prediction = prediction[:, :, : len(target_positions)]
 
-            # output format is (batch_size, prediction_length, n_quantiles)
-            # in case of quantile loss, the output n_quantiles = self.n_quantiles
-            # which is the length of a list of float. In case of MAE, MSE, etc.
-            # n_quantiles = 1 and it mimics the behavior of a point prediction.
-            # for multi-target forecasting, the output is a list of tensors.
+            # output format is (batch_size, prediction_length, n_targets)
             if len(target_positions) == 1:
                 prediction = prediction[..., 0, :]
             else:

pytorch_forecasting/models/timexer/_timexer_pkg.py

@@ -9,12 +9,12 @@ class TimeXer_pkg(_BasePtForecaster):
     _tags = {
         "info:name": "TimeXer",
         "info:compute": 3,
-        "info:pred_type": ["point", "quantile"],
+        "info:pred_type": ["point"],
         "info:y_type": ["numeric"],
         "authors": ["PranavBhatP"],
         "capability:exogenous": True,
         "capability:multivariate": True,
-        "capability:pred_int": True,
+        "capability:pred_int": False,
         "capability:flexible_history_length": True,
         "capability:cold_start": False,
     }

pytorch_forecasting/models/timexer/_timexer_pkg_v2.py

@@ -15,7 +15,7 @@ class TimeXer_pkg_v2(Base_pkg):
         "info:y_type": ["numeric"],
         "capability:exogenous": True,
         "capability:multivariate": True,
-        "capability:pred_int": True,
+        "capability:pred_int": False,
         "capability:flexible_history_length": False,
         "capability:cold_start": False,
     }

pytorch_forecasting/models/timexer/_timexer_v2.py

@@ -195,11 +195,6 @@ def _init_network(self):
         # float values.
         self.enc_in = self.enc_in or self.cont_dim
 
-        self.n_quantiles = None
-
-        if hasattr(self.loss, "quantiles") and self.loss.quantiles is not None:
-            self.n_quantiles = len(self.loss.quantiles)
-
         if self.hidden_size % self.n_heads != 0:
             raise ValueError(
                 f"hidden_size ({self.hidden_size}) must be divisible by n_heads ({self.n_heads}) "  # noqa: E501
@@ -259,7 +254,6 @@ def _init_network(self):
             self.head_nf,
             self.prediction_length,
             head_dropout=self.dropout,
-            n_quantiles=self.n_quantiles,
         )
 
     def _forecast(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:

pytorch_forecasting/models/timexer/sub_modules.py

@@ -259,17 +259,14 @@ class FlattenHead(nn.Module):
         Target window size.
     head_dropout : float, optional
         Dropout rate for the head. Defaults to 0.
-    n_quantiles : int, optional
-        Number of quantiles. Defaults to 1.
     """
 
-    def __init__(self, n_vars, nf, target_window, head_dropout=0, n_quantiles=1):
+    def __init__(self, n_vars, nf, target_window, head_dropout=0):
         super().__init__()
         self.n_vars = n_vars
         self.flatten = nn.Flatten(start_dim=-2)
-        self.n_quantiles = n_quantiles
 
-        self.linear = nn.Linear(nf, target_window * n_quantiles)
+        self.linear = nn.Linear(nf, target_window)
         self.dropout = nn.Dropout(head_dropout)
 
     def forward(self, x):
@@ -278,7 +275,7 @@ def forward(self, x):
         x = self.dropout(x)
 
         batch_size, n_vars = x.shape[0], x.shape[1]
-        x = x.reshape(batch_size, n_vars, -1, self.n_quantiles)
+        x = x.reshape(batch_size, n_vars, -1)
         return x