[ENH] Standardize output format for tslib v2 models (#1965)

Fixes #1964 
This PR standardizes the output format of `forward` of `tslib` models
to:
- [x]  3D tensors for single-target for `DLinear`
- [x] 3D tensors for single-target for `TimeXer`

Author: phoeenniixx

pytorch_forecasting/layers/_output/_flatten_head.py

@@ -38,8 +38,9 @@ def forward(self, x):
         x = self.flatten(x)
         x = self.linear(x)
         x = self.dropout(x)
+        x = x.permute(0, 2, 1)
 
         if self.n_quantiles is not None:
-            batch_size, n_vars = x.shape[0], x.shape[1]
-            x = x.reshape(batch_size, n_vars, -1, self.n_quantiles)
+            batch_size = x.shape[0]
+            x = x.reshape(batch_size, -1, self.n_quantiles)
         return x

pytorch_forecasting/models/dlinear/_dlinear_v2.py

@@ -239,22 +239,17 @@ def _reshape_output(self, output: torch.Tensor) -> torch.Tensor:
         Returns
         -------
         output: torch.Tensor
-            Reshaped tensor (batch_size, prediction_length, n_features, n_quantiles)
+            Reshaped tensor (batch_size, prediction_length, n_quantiles)
             or (batch_size, prediction_length, n_features) if n_quantiles is None.
         """
         if self.n_quantiles is not None:
-            batch_size, n_features = output.shape[0], output.shape[1]
+            batch_size = output.shape[0]
             output = output.reshape(
-                batch_size, n_features, self.prediction_length, self.n_quantiles
+                batch_size, self.prediction_length, self.n_quantiles
             )
-            output = output.permute(0, 2, 1, 3)  # (batch, time, features, quantiles)
         else:
             output = output.permute(0, 2, 1)  # (batch, time, features)
 
-        # univariate forecasting
-        if self.target_dim == 1 and output.shape[-1] == 1:
-            output = output.squeeze(-1)
-
         return output
 
     def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:

pytorch_forecasting/models/timexer/_timexer_v2.py

@@ -311,11 +311,6 @@ def _forecast(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
 
         dec_out = self.head(enc_out)
 
-        if self.n_quantiles is not None:
-            dec_out = dec_out.permute(0, 2, 1, 3)
-        else:
-            dec_out = dec_out.permute(0, 2, 1)
-
         return dec_out
 
     def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
@@ -330,10 +325,6 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         out = self._forecast(x)
         prediction = out[:, : self.prediction_length, :]
 
-        # check to see if the output shape is equal to number of targets
-        if prediction.size(2) != self.target_dim:
-            prediction = prediction[:, :, : self.target_dim]
-
         if "target_scale" in x:
             prediction = self.transform_output(prediction, x["target_scale"])
 

pytorch_forecasting/tests/test_all_estimators_v2.py

@@ -6,6 +6,7 @@
 import lightning.pytorch as pl
 from lightning.pytorch.callbacks import EarlyStopping
 from lightning.pytorch.loggers import TensorBoardLogger
+import torch
 import torch.nn as nn
 
 from pytorch_forecasting.tests.test_all_estimators import (
@@ -77,6 +78,19 @@ def _integration(
     test_outputs = trainer.test(net, dataloaders=test_dataloader)
     assert len(test_outputs) > 0
 
+    # todo: add the predict pipeline and make this test cleaner
+    x, y = next(iter(test_dataloader))
+    net.eval()
+    with torch.no_grad():
+        output = net(x)
+    net.train()
+    prediction = output["prediction"]
+    n_dims = prediction.ndim
+    assert n_dims == 3, (
+        f"Prediction output must be 3D, but got {n_dims}D tensor "
+        f"with shape {output.shape}"
+    )
+
     shutil.rmtree(tmp_path, ignore_errors=True)
 
 

tests/test_models/test_dlinear_v2.py

@@ -124,7 +124,7 @@ def test_quantile_loss_output(sample_dataset):
 
     assert "prediction" in output
     pred = output["prediction"]
-    assert pred.ndim == 4
+    assert pred.ndim == 3
     assert pred.shape[-1] == len(quantiles)
     assert pred.shape[1] == metadata["prediction_length"]
 

tests/test_models/test_timexer_v2.py

@@ -297,7 +297,7 @@ def test_quantile_predictions(basic_metadata):
         output = model(sample_input_data)
 
     predictions = output["prediction"]
-    assert predictions.shape == (batch_size, 8, 1, 3)
+    assert predictions.shape == (batch_size, 8, 3)
 
 
 def test_missing_history_target_handling(basic_metadata):