support broader losses and prameterize tests

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

Author: Faakhir30

pytorch_forecasting/metrics/nn_loss_adapter.py

@@ -1,85 +1,149 @@
+"""Adapter for native ``torch.nn`` loss modules in ptf-v2."""
+
+from __future__ import annotations
+
+import copy
+from typing import Literal
+
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
+
+_Mode = Literal["point", "class", "gaussian_nll"]
+_CLASS_LOSSES = (nn.CrossEntropyLoss, nn.NLLLoss)
+_GAUSSIAN_NLL_LOSSES = (nn.GaussianNLLLoss,)
 
 
 class NNLossAdapter(nn.Module):
-    """Adapter to use PyTorch nn losses in ptf-v2.
+    """Adapt a ``torch.nn`` loss module to the ptf-v2 loss API.
 
-    This class wraps a standard PyTorch loss (nn.Module) to handle the specific
+    Wraps a standard PyTorch loss (nn.Module) to handle the specific
     input formats used in pytorch-forecasting v2, such as (target, weight) tuples
     and multi-target list of tensors.
 
-    Args:
-        loss (nn.Module): The PyTorch loss to wrap.
+    The reshape mode is inferred automatically from the wrapped loss type:
+
+    * **point** — same-shape ``[B, T]`` after squeeze (default)
+    * **class** — logits ``[B, T, C]`` vs labels ``[B, T]``
+      (``CrossEntropyLoss``, ``NLLLoss``)
+    * **gaussian_nll** — mean/var head ``[B, T, 2]`` (``GaussianNLLLoss``)
+
+    Parameters
+    ----------
+    loss :
+        Native PyTorch loss, e.g. ``nn.MSELoss()``.
     """
 
     def __init__(self, loss: nn.Module):
         super().__init__()
-        self._loss = loss
+        # deepcopy so we never mutate the caller's loss instance
+        self._loss = copy.deepcopy(loss)
+        self._reduction = getattr(loss, "reduction", "mean")
+        self._mode = self._infer_mode(self._loss)
+
+    @staticmethod
+    def _infer_mode(loss: nn.Module) -> _Mode:
+        if isinstance(loss, _CLASS_LOSSES):
+            return "class"
+        if isinstance(loss, _GAUSSIAN_NLL_LOSSES):
+            return "gaussian_nll"
+        return "point"
 
     def forward(
         self,
         y_pred: torch.Tensor | list[torch.Tensor],
-        y_actual: torch.Tensor | tuple[torch.Tensor | list[torch.Tensor], torch.Tensor],
+        y_actual: torch.Tensor
+        | list[torch.Tensor]
+        | tuple[torch.Tensor | list[torch.Tensor], torch.Tensor],
     ) -> torch.Tensor:
         """
         Forward pass of the adapter.
 
-        Args:
-            y_pred (torch.Tensor | list[torch.Tensor]): Model predictions.
-                Expected to be [B, T, N] for multi-target or
-                [B, T, 1] / [B, T] for single target.
-            y_actual (torch.Tensor | tuple): Actual values and optionally weights.
-                Can be a tensor, or a tuple (target, weight), where target
-                can be a list of tensors.
+        Parameters
+        ----------
+        y_pred :
+            Model predictions.
+
+            * point: ``[B, T, 1]`` / ``[B, T]``, or ``[B, T, N]`` for multi-target
+            * class: ``[B, T, C]`` logits
+            * gaussian_nll: ``[B, T, 2]`` as ``(mean, raw_variance)``
+        y_actual :
+            Targets, optionally as ``(target, weight)``. Multi-target uses a
+            list of ``[B, T]`` tensors (point mode only).
 
-        Returns:
-            torch.Tensor: The computed and reduced loss.
+        Returns
+        -------
+        torch.Tensor
+            Scalar (or unreduced) loss.
         """
-        # Handle y_actual as (target, weight) or just target
-        if isinstance(y_actual, (list, tuple)) and not isinstance(
-            y_actual, torch.Tensor
-        ):
-            if len(y_actual) == 2:
-                target, weight = y_actual
-            else:
-                target = y_actual[0]
-                weight = None
-        else:
-            target, weight = y_actual, None
+        target, weight = self._unpack_y_actual(y_actual)
+        mode = self._mode
 
+        # multi-target scenario
         if isinstance(target, list):
-            # Multi-target scenario
+            if mode != "point":
+                raise ValueError(
+                    f"Error in NNLossAdapter: Multi-target lists are only supported"
+                    f"for point losses, got {self._loss.__class__.__name__!r}."
+                )
             if not isinstance(y_pred, torch.Tensor):
                 raise ValueError(
                     f"NNLossAdapter expected y_pred to be a torch.Tensor for "
                     f"multi-target, but got {type(y_pred)}. Standard multi-target "
                     f"in ptf-v2 expects y_pred of shape [B, T, N]."
                 )
-
             # y_pred is [B, T, N], split along last dimension
-            y_preds = y_pred.split(1, dim=-1)
-            y_preds = [yp.squeeze(-1) for yp in y_preds]
-
+            y_preds = [yp.squeeze(-1) for yp in y_pred.split(1, dim=-1)]
             if len(y_preds) != len(target):
                 raise ValueError(
                     f"Number of predictions ({len(y_preds)}) does not match "
                     f"number of targets ({len(target)})."
                 )
 
-            total_loss = torch.tensor(0.0, device=y_pred.device)
+            total_loss = torch.tensor(0.0, device=y_pred.device, dtype=y_pred.dtype)
             for yp, t in zip(y_preds, target):
-                total_loss = total_loss + self._compute_loss(yp, t, weight)
+                total_loss = total_loss + self._compute_loss(yp, t, weight, mode)
             return total_loss
-        else:
-            # Single target scenario
-            if isinstance(y_pred, list):
-                # Error if list of predictions but single tensor target
-                raise ValueError(
-                    "NNLossAdapter does not support list of predictions "
-                    "with single target tensor."
-                )
 
+        # single-target scenario
+        if isinstance(y_pred, list):
+            raise ValueError(
+                "NNLossAdapter does not support list of predictions "
+                "with single target tensor."
+            )
+
+        y_pred, target = self._prepare_inputs(y_pred, target, mode)
+        return self._compute_loss(y_pred, target, weight, mode)
+
+    @staticmethod
+    def _unpack_y_actual(
+        y_actual: torch.Tensor | list[torch.Tensor] | tuple,
+    ) -> tuple[torch.Tensor | list[torch.Tensor], torch.Tensor | None]:
+        if (
+            isinstance(y_actual, tuple)
+            and len(y_actual) == 2
+            and torch.is_tensor(y_actual[1])
+        ):
+            return y_actual[0], y_actual[1]
+        # also allow (target, None) or len-2 list/tuple without weight tensor
+        if isinstance(y_actual, (list, tuple)) and not isinstance(
+            y_actual, torch.Tensor
+        ):
+            if len(y_actual) == 2 and (
+                y_actual[1] is None or torch.is_tensor(y_actual[1])
+            ):
+                return y_actual[0], y_actual[1]
+            if len(y_actual) == 1:
+                return y_actual[0], None
+        return y_actual, None
+
+    def _prepare_inputs(
+        self,
+        y_pred: torch.Tensor,
+        target: torch.Tensor,
+        mode: _Mode,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if mode == "point":
             if y_pred.ndim == 3:
                 if y_pred.size(-1) != 1:
                     raise ValueError(
@@ -89,25 +153,73 @@ def forward(
                         "For multi-horizon losses, use a ptf metrics loss instead."
                     )
                 y_pred = y_pred.squeeze(-1)
+            return y_pred, target
 
-            return self._compute_loss(y_pred, target, weight)
+        if mode == "class":
+            if y_pred.ndim != 3:
+                raise ValueError(
+                    "Classification losses expect logits of shape "
+                    f"(batch, time, classes), got {tuple(y_pred.shape)}."
+                )
+            if target.ndim != 2:
+                raise ValueError(
+                    "Classification targets must have shape (batch, time), "
+                    f"got {tuple(target.shape)}."
+                )
+            return y_pred.reshape(-1, y_pred.size(-1)), target.reshape(-1).long()
+
+        # gaussian_nll
+        if y_pred.ndim != 3 or y_pred.size(-1) != 2:
+            raise ValueError(
+                "GaussianNLLLoss expects predictions of shape "
+                f"(batch, time, 2) as (mean, raw_variance); got {tuple(y_pred.shape)}."
+            )
+        if target.ndim != 2:
+            raise ValueError(
+                "GaussianNLL targets must have shape (batch, time), "
+                f"got {tuple(target.shape)}."
+            )
+        return y_pred, target
+
+    def _call_loss(
+        self,
+        y_pred: torch.Tensor,
+        target: torch.Tensor,
+        mode: _Mode,
+    ) -> torch.Tensor:
+        if mode == "gaussian_nll":
+            mean = y_pred[..., 0]
+            var = F.softplus(y_pred[..., 1]) + 1e-6
+            return self._loss(mean, target, var)
+        if mode == "class" and isinstance(self._loss, nn.NLLLoss):
+            # NLLLoss expects log-probabilities
+            return self._loss(F.log_softmax(y_pred, dim=-1), target)
+        return self._loss(y_pred, target)
 
     def _compute_loss(
-        self, y_pred: torch.Tensor, target: torch.Tensor, weight: torch.Tensor | None
+        self,
+        y_pred: torch.Tensor,
+        target: torch.Tensor,
+        weight: torch.Tensor | None,
+        mode: _Mode,
     ) -> torch.Tensor:
-        """
-        Compute the loss for a single target, applying weights if provided.
-        """
         if weight is None:
-            return self._loss(y_pred, target)
+            return self._call_loss(y_pred, target, mode)
 
-        # Handle weighting
-        old_reduction = getattr(self._loss, "reduction", "mean")
-        self._loss.reduction = "none"
+        old_reduction = getattr(self._loss, "reduction", None)
+        if old_reduction is not None:
+            self._loss.reduction = "none"
         try:
-            loss = self._loss(y_pred, target)
+            loss = self._call_loss(y_pred, target, mode)
         finally:
-            self._loss.reduction = old_reduction
+            if old_reduction is not None:
+                self._loss.reduction = old_reduction
+
+        # class mode flattens to (B*T,); flatten matching weights
+        if mode == "class" and weight is not None:
+            weight = weight.reshape(-1)
+        elif mode == "gaussian_nll" and weight is not None and weight.ndim == 2:
+            pass  # already [B, T], matches per-element loss
 
         # Ensure weight has same dimensions as loss for multiplication
         if weight.ndim < loss.ndim:
@@ -122,23 +234,29 @@ def _compute_loss(
             return weighted_loss.sum() / weight.sum()
         elif old_reduction == "sum":
             return weighted_loss.sum()
-        else:
-            # 'none' or others
-            return weighted_loss
+        # 'none' or others
+        return weighted_loss
 
     def to_prediction(self, y_pred: torch.Tensor, **kwargs) -> torch.Tensor:
         """
         Convert network prediction into a point prediction.
         """
-        if y_pred.ndim == 3:
-            if y_pred.size(-1) == 1:
-                return y_pred.squeeze(-1)
+        del kwargs
+        mode = self._mode
+        if mode == "class" and y_pred.ndim == 3:
+            return y_pred.argmax(dim=-1)
+        if mode == "gaussian_nll" and y_pred.ndim == 3 and y_pred.size(-1) == 2:
+            return y_pred[..., 0]
+        if y_pred.ndim == 3 and y_pred.size(-1) == 1:
+            return y_pred.squeeze(-1)
         return y_pred
 
     def to_quantiles(self, y_pred: torch.Tensor, **kwargs) -> torch.Tensor:
         """
         Convert network prediction into a quantile prediction.
         """
-        if y_pred.ndim == 2:
-            return y_pred.unsqueeze(-1)
-        return y_pred
+        del kwargs
+        point = self.to_prediction(y_pred)
+        if point.ndim == 2:
+            return point.unsqueeze(-1)
+        return point

pytorch_forecasting/models/temporal_fusion_transformer/_tft_pkg_v2.py

@@ -59,6 +59,12 @@ def get_test_train_params(cls):
                 hidden_size=16,
                 attention_head_size=4,
             ),
+            dict(
+                loss=nn.GaussianNLLLoss(),
+                output_size=2,
+                hidden_size=16,
+                attention_head_size=2,
+            ),
             dict(datamodule_cfg=dict(max_encoder_length=5, max_prediction_length=3)),
             dict(
                 hidden_size=24,