content: __all__ = ['BasePointLoss', 'MAE', 'MSE', 'RMSE', 'MAPE', 'SMAPE', 'MASE', 'relMSE', 'QuantileLoss', 'MQLoss', 'QuantileLayer',
             'IQLoss', 'DistributionLoss', 'PMM', 'GMM', 'NBMM', 'HuberLoss', 'TukeyLoss', 'HuberQLoss', 'HuberMQLoss',
             'HuberIQLoss', 'Accuracy', 'sCRPS']
  
  
  from functools import partial
  from typing import List, Optional, Tuple, Union
  
  import numpy as np
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.distributions import (
      AffineTransform,
      Bernoulli,
      Beta,
      Categorical,
      Distribution,
      Gamma,
      MixtureSameFamily,
      NegativeBinomial,
      Normal,
      Poisson,
      StudentT,
      TransformedDistribution,
      constraints,
  )
  
  
  def _divide_no_nan(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
      """
      Auxiliary funtion to handle divide by 0
      """
      div = a / b
      return torch.nan_to_num(div, nan=0.0, posinf=0.0, neginf=0.0)
  
  
  def _weighted_mean(losses, weights):
      """
      Compute weighted mean of losses per datapoint.
      """
      return _divide_no_nan(torch.sum(losses * weights), torch.sum(weights))
  
  
  class BasePointLoss(torch.nn.Module):
      """Base class for point loss functions.
  
      Args:
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
          outputsize_multiplier (Optional[int]): Multiplier for the output size. Defaults to None.
          output_names (Optional[List[str]]): Names of the outputs. Defaults to None.
      """
  
      def __init__(
          self, horizon_weight=None, outputsize_multiplier=None, output_names=None
      ):
          super(BasePointLoss, self).__init__()
          if horizon_weight is not None:
              horizon_weight = torch.Tensor(horizon_weight.flatten())
          self.horizon_weight = horizon_weight
          self.outputsize_multiplier = outputsize_multiplier
          self.output_names = output_names
          self.is_distribution_output = False
  
      def domain_map(self, y_hat: torch.Tensor):
          """Domain mapping for predicted values.
  
          Args:
              y_hat (torch.Tensor): Predicted values tensor.
                  - Univariate: [B, H, 1]
                  - Multivariate: [B, H, N]
  
          Returns:
              torch.Tensor: Mapped values tensor with shape [B, H, N].
          """
          return y_hat
  
      def _compute_weights(self, y, mask):
          """Compute final weights for each datapoint based on all weights and masks.
  
          Set horizon_weight to a ones[H] tensor if not set.
          If set, check that it has the same length as the horizon in x.
  
          Args:
              y (torch.Tensor): Target values tensor.
              mask (torch.Tensor, optional): Mask tensor specifying datapoints to consider.
  
          Returns:
              torch.Tensor: Final weights tensor for each datapoint.
          """
          if mask is None:
              mask = torch.ones_like(y)
  
          if self.horizon_weight is None:
              weights = torch.ones_like(mask)
          else:
              assert mask.shape[1] == len(
                  self.horizon_weight
              ), "horizon_weight must have same length as Y"
              weights = self.horizon_weight.clone()
              weights = weights[None, :, None].to(mask.device)
              weights = torch.ones_like(mask, device=mask.device) * weights
  
          return weights * mask
  
  
  class MAE(BasePointLoss):
      r"""Mean Absolute Error.
  
      Calculates Mean Absolute Error between `y` and `y_hat`. MAE measures the relative prediction
      accuracy of a forecasting method by calculating the deviation of the prediction and the true
      value at a given time and averages these devations over the length of the series.
  
      ```math
      \mathrm{MAE}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} |y_{\tau} - \hat{y}_{\tau}|
      ```
  
      Args:
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
      """
  
      def __init__(self, horizon_weight=None):
          super(MAE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          mask: Union[torch.Tensor, None] = None,
          y_insample: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """Calculate Mean Absolute Error between actual and predicted values.
  
          Args:
              y (torch.Tensor): Actual values.
              y_hat (torch.Tensor): Predicted values.
              mask (Union[torch.Tensor, None], optional): Specifies datapoints to consider in loss. Defaults to None.
              y_insample (Union[torch.Tensor, None], optional): Actual insample values. Defaults to None.
  
          Returns:
              torch.Tensor: MAE (single value).
          """
          losses = torch.abs(y - y_hat)
          weights = self._compute_weights(y=y, mask=mask)
          return _weighted_mean(losses=losses, weights=weights)
  
  
  class MSE(BasePointLoss):
      r"""Mean Squared Error.
  
      Calculates Mean Squared Error between `y` and `y_hat`. MSE measures the relative prediction
      accuracy of a forecasting method by calculating the squared deviation of the prediction and the true
      value at a given time, and averages these devations over the length of the series.
  
      ```math
      \mathrm{MSE}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} (y_{\tau} - \hat{y}_{\tau})^{2} 
      ```
  
      Args:
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
      """
  
      def __init__(self, horizon_weight=None):
          super(MSE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          y_insample: Union[torch.Tensor, None] = None,
          mask: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """Calculate Mean Squared Error between actual and predicted values.
  
          Args:
              y (torch.Tensor): Actual values.
              y_hat (torch.Tensor): Predicted values.
              y_insample (Union[torch.Tensor, None], optional): Actual insample values. Defaults to None.
              mask (Union[torch.Tensor, None], optional): Specifies datapoints to consider in loss. Defaults to None.
  
          Returns:
              torch.Tensor: MSE (single value).
          """
          losses = (y - y_hat) ** 2
          weights = self._compute_weights(y=y, mask=mask)
          return _weighted_mean(losses=losses, weights=weights)
  
  
  class RMSE(BasePointLoss):
      r"""Root Mean Squared Error.
  
      Calculates Root Mean Squared Error between `y` and `y_hat`. RMSE measures the relative prediction
      accuracy of a forecasting method by calculating the squared deviation of the prediction and the observed value at
      a given time and averages these devations over the length of the series.
      Finally the RMSE will be in the same scale as the original time series so its comparison with other
      series is possible only if they share a common scale. RMSE has a direct connection to the L2 norm.
  
      ```math
      \mathrm{RMSE}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}_{\tau}) = \sqrt{\frac{1}{H} \sum^{t+H}_{\tau=t+1} (y_{\tau} - \hat{y}_{\tau})^{2}}
      ```
  
      Args:
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
      """
  
      def __init__(self, horizon_weight=None):
          super(RMSE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          mask: Union[torch.Tensor, None] = None,
          y_insample: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """
          Args:
              y: Tensor, Actual values.
              y_hat: Tensor, Predicted values.
              mask: Tensor, Specifies datapoints to consider in loss.
  
          Returns:
              rmse: Tensor (single value).
          """
          losses = (y - y_hat) ** 2
          weights = self._compute_weights(y=y, mask=mask)
          losses = _weighted_mean(losses=losses, weights=weights)
          return torch.sqrt(losses)
  
  
  class MAPE(BasePointLoss):
      r"""Mean Absolute Percentage Error
  
      Calculates Mean Absolute Percentage Error  between
      `y` and `y_hat`. MAPE measures the relative prediction
      accuracy of a forecasting method by calculating the percentual deviation
      of the prediction and the observed value at a given time and
      averages these devations over the length of the series.
      The closer to zero an observed value is, the higher penalty MAPE loss
      assigns to the corresponding error.
  
      ```math
      \mathrm{MAPE}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} \frac{|y_{\tau}-\hat{y}_{\tau}|}{|y_{\tau}|}
      ```
  
  
      Args:
          horizon_weight: Tensor of size h, weight for each timestamp of the forecasting window.
  
      References:
          - [Makridakis S., "Accuracy measures: theoretical and practical concerns".](https://www.sciencedirect.com/science/article/pii/0169207093900793)
      """
  
      def __init__(self, horizon_weight=None):
          super(MAPE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          y_insample: Union[torch.Tensor, None] = None,
          mask: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """
          Args:
              y: Tensor, Actual values.
              y_hat: Tensor, Predicted values.
              mask: Tensor, Specifies date stamps per serie to consider in loss.
  
          Returns:
              mape: Tensor (single value).
          """
          scale = _divide_no_nan(torch.ones_like(y, device=y.device), torch.abs(y))
          losses = torch.abs(y - y_hat) * scale
          weights = self._compute_weights(y=y, mask=mask)
          mape = _weighted_mean(losses=losses, weights=weights)
          return mape
  
  
  class SMAPE(BasePointLoss):
      r"""Symmetric Mean Absolute Percentage Error
  
      Calculates Symmetric Mean Absolute Percentage Error between
      `y` and `y_hat`. SMAPE measures the relative prediction
      accuracy of a forecasting method by calculating the relative deviation
      of the prediction and the observed value scaled by the sum of the
      absolute values for the prediction and observed value at a
      given time, then averages these devations over the length
      of the series. This allows the SMAPE to have bounds between
      0% and 200% which is desireble compared to normal MAPE that
      may be undetermined when the target is zero.
  
      ```math
      \mathrm{sMAPE}_{2}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} \frac{|y_{\tau}-\hat{y}_{\tau}|}{|y_{\tau}|+|\hat{y}_{\tau}|}
      ```
  
  
      Args:
          horizon_weight: Tensor of size h, weight for each timestamp of the forecasting window.
  
      References:
          - [Makridakis S., "Accuracy measures: theoretical and practical concerns".](https://www.sciencedirect.com/science/article/pii/0169207093900793)
      """
  
      def __init__(self, horizon_weight=None):
          super(SMAPE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          mask: Union[torch.Tensor, None] = None,
          y_insample: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """
          Args:
              y: Tensor, Actual values.
              y_hat: Tensor, Predicted values.
              mask: Tensor, Specifies date stamps per serie to consider in loss.
  
          Returns:
              smape: Tensor (single value).
          """
          delta_y = torch.abs((y - y_hat))
          scale = torch.abs(y) + torch.abs(y_hat)
          losses = _divide_no_nan(delta_y, scale)
          weights = self._compute_weights(y=y, mask=mask)
          return 2 * _weighted_mean(losses=losses, weights=weights)
  
  
  class MASE(BasePointLoss):
      r"""Mean Absolute Scaled Error
      Calculates the Mean Absolute Scaled Error between
      `y` and `y_hat`. MASE measures the relative prediction
      accuracy of a forecasting method by comparinng the mean absolute errors
      of the prediction and the observed value against the mean
      absolute errors of the seasonal naive model.
      The MASE partially composed the Overall Weighted Average (OWA),
      used in the M4 Competition.
  
      ```math
      \mathrm{MASE}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}_{\tau}, \mathbf{\hat{y}}^{season}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} \frac{|y_{\tau}-\hat{y}_{\tau}|}{\mathrm{MAE}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}^{season}_{\tau})}
      ```
  
  
      Args:
          seasonality: Int. Main frequency of the time series; Hourly 24,  Daily 7, Weekly 52, Monthly 12, Quarterly 4, Yearly 1.
          horizon_weight: Tensor of size h, weight for each timestamp of the forecasting window.
  
      References:
          [Rob J. Hyndman, & Koehler, A. B. "Another look at measures of forecast accuracy".](https://www.sciencedirect.com/science/article/pii/S0169207006000239)
          [Spyros Makridakis, Evangelos Spiliotis, Vassilios Assimakopoulos, "The M4 Competition: 100,000 time series and 61 forecasting methods".](https://www.sciencedirect.com/science/article/pii/S0169207019301128)
      """
  
      def __init__(self, seasonality: int, horizon_weight=None):
          super(MASE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
          self.seasonality = seasonality
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          y_insample: torch.Tensor,
          mask: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """
          Args:
              y: Tensor (batch_size, output_size), Actual values.
              y_hat: Tensor (batch_size, output_size)), Predicted values.
              y_insample: Tensor (batch_size, input_size), Actual insample values.
              mask: Tensor, Specifies date stamps per serie to consider in loss.
  
          Returns:
              mase: Tensor (single value).
          """
          delta_y = torch.abs(y - y_hat)
          scale = torch.mean(
              torch.abs(
                  y_insample[:, self.seasonality :] - y_insample[:, : -self.seasonality]
              ),
              axis=1,
          )
          losses = _divide_no_nan(delta_y, scale[:, None, None])
          weights = self._compute_weights(y=y, mask=mask)
          return _weighted_mean(losses=losses, weights=weights)
  
  
  class relMSE(BasePointLoss):
      r"""Relative Mean Squared Error
      Computes Relative Mean Squared Error (relMSE), as proposed by Hyndman & Koehler (2006)
      as an alternative to percentage errors, to avoid measure unstability.
  
      ```math
      \mathrm{relMSE}(\mathbf{y}, \mathbf{\hat{y}}, \mathbf{\hat{y}}^{benchmark}) =
      \frac{\mathrm{MSE}(\mathbf{y}, \mathbf{\hat{y}})}{\mathrm{MSE}(\mathbf{y}, \mathbf{\hat{y}}^{benchmark})}
      ```
  
      Args:
          y_train: Numpy array, deprecated.
          horizon_weight: Tensor of size h, weight for each timestamp of the forecasting window.
  
      References:
          - [Hyndman, R. J and Koehler, A. B. (2006). "Another look at measures of forecast accuracy", International Journal of Forecasting, Volume 22, Issue 4.](https://www.sciencedirect.com/science/article/pii/S0169207006000239)
          - [Kin G. Olivares, O. Nganba Meetei, Ruijun Ma, Rohan Reddy, Mengfei Cao, Lee Dicker. "Probabilistic Hierarchical Forecasting with Deep Poisson Mixtures. Submitted to the International Journal Forecasting, Working paper available at arxiv.](https://arxiv.org/pdf/2110.13179.pdf)
      """
  
      def __init__(self, y_train=None, horizon_weight=None):
          super(relMSE, self).__init__(
              horizon_weight=horizon_weight, outputsize_multiplier=1, output_names=[""]
          )
          if y_train is not None:
              raise DeprecationWarning("y_train will be deprecated in a future release.")
          self.mse = MSE(horizon_weight=horizon_weight)
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          y_benchmark: torch.Tensor,
          mask: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """
          Args:
              y: Tensor (batch_size, output_size), Actual values.
              y_hat: Tensor (batch_size, output_size)), Predicted values.
              y_benchmark: Tensor (batch_size, output_size), Benchmark predicted values.
              mask: Tensor, Specifies date stamps per serie to consider in loss.
  
          Returns:
              relMSE: Tensor (single value).
          """
          norm = self.mse(y=y, y_hat=y_benchmark, mask=mask)  # Already weighted
          norm = norm + 1e-5  # Numerical stability
          loss = self.mse(y=y, y_hat=y_hat, mask=mask)  # Already weighted
          loss = _divide_no_nan(loss, norm)
          return loss
  
  
  class QuantileLoss(BasePointLoss):
      r"""Quantile Loss.
  
      Computes the quantile loss between `y` and `y_hat`.
      QL measures the deviation of a quantile forecast.
      By weighting the absolute deviation in a non symmetric way, the
      loss pays more attention to under or over estimation.
      A common value for q is 0.5 for the deviation from the median (Pinball loss).
  
      ```math
      \mathrm{QL}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}^{(q)}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} \Big( (1-q)\,( \hat{y}^{(q)}_{\tau} - y_{\tau} )_{+} + q\,( y_{\tau} - \hat{y}^{(q)}_{\tau} )_{+} \Big)
      ```
  
      Args:
          q (float): Between 0 and 1. The slope of the quantile loss, in the context of quantile regression, the q determines the conditional quantile level.
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
  
      References:
          [Roger Koenker and Gilbert Bassett, Jr., "Regression Quantiles".](https://www.jstor.org/stable/1913643)
      """
  
      def __init__(self, q, horizon_weight=None):
          super(QuantileLoss, self).__init__(
              horizon_weight=horizon_weight,
              outputsize_multiplier=1,
              output_names=[f"_ql{q}"],
          )
          self.q = q
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          y_insample: Union[torch.Tensor, None] = None,
          mask: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          r"""Calculate quantile loss between actual and predicted values.
  
          Args:
              y (torch.Tensor): Actual values.
              y_hat (torch.Tensor): Predicted values.
              y_insample (Union[torch.Tensor, None], optional): Actual insample values. Defaults to None.
              mask (Union[torch.Tensor, None], optional): Specifies datapoints to consider in loss. Defaults to None.
  
          Returns:
              torch.Tensor: Quantile loss (single value).
          """
          delta_y = y - y_hat
          losses = torch.max(torch.mul(self.q, delta_y), torch.mul((self.q - 1), delta_y))
          weights = self._compute_weights(y=y, mask=mask)
          return _weighted_mean(losses=losses, weights=weights)
  
  
  def level_to_outputs(level):
      qs = sum([[50 - l / 2, 50 + l / 2] for l in level], [])
      output_names = sum([[f"-lo-{l}", f"-hi-{l}"] for l in level], [])
  
      sort_idx = np.argsort(qs)
      quantiles = np.array(qs)[sort_idx]
  
      # Add default median
      quantiles = np.concatenate([np.array([50]), quantiles])
      quantiles = torch.Tensor(quantiles) / 100
      output_names = list(np.array(output_names)[sort_idx])
      output_names.insert(0, "-median")
  
      return quantiles, output_names
  
  
  def quantiles_to_outputs(quantiles):
      output_names = []
      for q in quantiles:
          if q < 0.50:
              output_names.append(f"-lo-{np.round(100-200*q,2)}")
          elif q > 0.50:
              output_names.append(f"-hi-{np.round(100-200*(1-q),2)}")
          else:
              output_names.append("-median")
      return quantiles, output_names
  
  
  class MQLoss(BasePointLoss):
      r"""Multi-Quantile loss
  
      Calculates the Multi-Quantile loss (MQL) between `y` and `y_hat`.
      MQL calculates the average multi-quantile Loss for
      a given set of quantiles, based on the absolute
      difference between predicted quantiles and observed values.
  
      ```math
      \mathrm{MQL}(\mathbf{y}_{\tau},[\mathbf{\hat{y}}^{(q_{1})}_{\tau}, ... ,\hat{y}^{(q_{n})}_{\tau}]) = \frac{1}{n} \sum_{q_{i}} \mathrm{QL}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}^{(q_{i})}_{\tau})
      ```
  
      The limit behavior of MQL allows to measure the accuracy
      of a full predictive distribution $\mathbf{\hat{F}}_{\tau}$ with
      the continuous ranked probability score (CRPS). This can be achieved
      through a numerical integration technique, that discretizes the quantiles
      and treats the CRPS integral with a left Riemann approximation, averaging over
      uniformly distanced quantiles.
  
      ```math
      \mathrm{CRPS}(y_{\tau}, \mathbf{\hat{F}}_{\tau}) = \int^{1}_{0} \mathrm{QL}(y_{\tau}, \hat{y}^{(q)}_{\tau}) dq
      ```
  
      Args:
          level (List[int], optional): Probability levels for prediction intervals. Defaults to [80, 90].
          quantiles (Optional[List[float]]): Alternative to level, quantiles to estimate from y distribution. Defaults to None.
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
  
      References:
          [Roger Koenker and Gilbert Bassett, Jr., "Regression Quantiles".](https://www.jstor.org/stable/1913643)
          [James E. Matheson and Robert L. Winkler, "Scoring Rules for Continuous Probability Distributions".](https://www.jstor.org/stable/2629907)
      """
  
      def __init__(self, level=[80, 90], quantiles=None, horizon_weight=None):
  
          qs, output_names = level_to_outputs(level)
          qs = torch.Tensor(qs)
          # Transform quantiles to homogeneous output names
          if quantiles is not None:
              _, output_names = quantiles_to_outputs(quantiles)
              qs = torch.Tensor(quantiles)
  
          super(MQLoss, self).__init__(
              horizon_weight=horizon_weight,
              outputsize_multiplier=len(qs),
              output_names=output_names,
          )
  
          self.quantiles = torch.nn.Parameter(qs, requires_grad=False)
  
      def domain_map(self, y_hat: torch.Tensor):
          """Reshapes input tensor to match the expected output format.
  
          Args:
              y_hat (torch.Tensor): Input tensor.
                  - Univariate: [B, H, 1 * Q]
                  - Multivariate: [B, H, N * Q]
  
          Returns:
              torch.Tensor: Reshaped tensor with shape [B, H, N, Q].
          """
          output = y_hat.reshape(
              y_hat.shape[0], y_hat.shape[1], -1, self.outputsize_multiplier
          )
  
          return output
  
      def _compute_weights(self, y, mask):
          """Compute final weights for each datapoint based on all weights and masks.
  
          Set horizon_weight to a ones[H] tensor if not set.
          If set, check that it has the same length as the horizon in x.
  
          Args:
              y: Tensor with shape [B, h, N, 1].
              mask: Tensor with shape [B, h, N, 1].
          """
  
          if self.horizon_weight is None:
              weights = torch.ones_like(mask)
          else:
              assert mask.shape[1] == len(
                  self.horizon_weight
              ), "horizon_weight must have same length as Y"
              weights = self.horizon_weight.clone()
              weights = weights[None, :, None, None]
              weights = weights.to(mask.device)
              weights = torch.ones_like(mask, device=mask.device) * weights
  
          return weights * mask
  
      def __call__(
          self,
          y: torch.Tensor,
          y_hat: torch.Tensor,
          y_insample: Union[torch.Tensor, None] = None,
          mask: Union[torch.Tensor, None] = None,
      ) -> torch.Tensor:
          """Computes the multi-quantile loss.
  
          Args:
              y (torch.Tensor): Actual values.
              y_hat (torch.Tensor): Predicted values.
              y_insample (Union[torch.Tensor, None], optional): In-sample values. Defaults to None.
              mask (Union[torch.Tensor, None], optional): Specifies date stamps per serie to consider in loss. Defaults to None.
  
          Returns:
              torch.Tensor: Multi-quantile loss (single value).
          """
          # [B, h, N] -> [B, h, N, 1]
          if y_hat.ndim == 3:
              y_hat = y_hat.unsqueeze(-1)
  
          y = y.unsqueeze(-1)
          if mask is not None:
              mask = mask.unsqueeze(-1)
          else:
              mask = torch.ones_like(y, device=y.device)
  
          error = y_hat - y
  
          sq = torch.maximum(-error, torch.zeros_like(error))
          s1_q = torch.maximum(error, torch.zeros_like(error))
  
          quantiles = self.quantiles[None, None, None, :]
          losses = (1 / len(quantiles)) * (quantiles * sq + (1 - quantiles) * s1_q)
          weights = self._compute_weights(y=losses, mask=mask)  # Use losses for extra dim
  
          return _weighted_mean(losses=losses, weights=weights)
  
  
  class QuantileLayer(nn.Module):
      """Implicit Quantile Layer from the paper IQN for Distributional Reinforcement Learning.
  
      Code from GluonTS: https://github.com/awslabs/gluonts/blob/61133ef6e2d88177b32ace4afc6843ab9a7bc8cd/src/gluonts/torch/distributions/implicit_quantile_network.py
  
      References:
          Dabney et al. 2018. https://arxiv.org/abs/1806.06923
      """
  
      def __init__(self, num_output: int, cos_embedding_dim: int = 128):
          super().__init__()
  
          self.output_layer = nn.Sequential(
              nn.Linear(cos_embedding_dim, cos_embedding_dim),
              nn.PReLU(),
              nn.Linear(cos_embedding_dim, num_output),
          )
  
          self.register_buffer("integers", torch.arange(0, cos_embedding_dim))
  
      def forward(self, tau: torch.Tensor) -> torch.Tensor:
          cos_emb_tau = torch.cos(tau * self.integers * torch.pi)
          return self.output_layer(cos_emb_tau)
  
  
  class IQLoss(QuantileLoss):
      r"""Implicit Quantile Loss.
  
      Computes the quantile loss between `y` and `y_hat`, with the quantile `q` provided as an input to the network.
      IQL measures the deviation of a quantile forecast.
      By weighting the absolute deviation in a non symmetric way, the
      loss pays more attention to under or over estimation.
  
      ```math
      \mathrm{QL}(\mathbf{y}_{\tau}, \mathbf{\hat{y}}^{(q)}_{\tau}) = \frac{1}{H} \sum^{t+H}_{\tau=t+1} \Big( (1-q)\,( \hat{y}^{(q)}_{\tau} - y_{\tau} )_{+} + q\,( y_{\tau} - \hat{y}^{(q)}_{\tau} )_{+} \Big)
      ```
  
      Args:
          cos_embedding_dim (int, optional): Cosine embedding dimension. Defaults to 64.
          concentration0 (float, optional): Beta distribution concentration parameter. Defaults to 1.0.
          concentration1 (float, optional): Beta distribution concentration parameter. Defaults to 1.0.
          horizon_weight (Optional[torch.Tensor]): Tensor of size h, weight for each timestamp of the forecasting window. Defaults to None.
  
      References:
          Gouttes, Adèle, Kashif Rasul, Mateusz Koren, Johannes Stephan, and Tofigh Naghibi, "Probabilistic Time Series Forecasting with Implicit Quantile Networks". http://arxiv.org/abs/2107.03743
      """
  
      def __init__(
          self,
          cos_embedding_dim=64,
          concentration0=1.0,
          concentration1=1.0,
          horizon_weight=None,
      ):
          self.update_quantile()
          super(IQLoss, self).__init__(q=self.q, horizon_weight=horizon_weight)
  
          self.cos_embedding_dim = cos_embedding_dim
          self.concentration0 = concentration0
          self.concentration1 = concentration1
          self.has_sampled = False
          self.has_predicted = False
  
          self.quantile_layer = QuantileLayer(
              num_output=1, cos_embedding_dim=self.cos_embedding_dim
          )
          self.output_layer = nn.Sequential(nn.Linear(1, 1), nn.PReLU())
  
      def _sample_quantiles(self, sample_size, device):
          if not self.has_sampled:
              self._init_sampling_distribution(device)
  
          quantiles = self.sampling_distr.sample(sample_size)
          self.q = quantiles.squeeze(-1)
          self.has_sampled = True
          self.has_predicted = False
  
          return quantiles
  
      def _init_sampling_distribution(self, device):
          concentration0 = torch.tensor(
              [self.concentration0], device=device, dtype=torch.float32
          )
          concentration1 = torch.tensor(
              [self.concentration1], device=device, dtype=torch.float32
          )
          self.sampling_distr = Beta(
              concentration0=concentration0, concentration1=concentration1
          )
  
      def update_quantile(self, q: List[float] = [0.5]):
          self.q = q[0]
          self.output_names = [f"_ql{q[0]}"]
          self.has_predicted = True
  
      def domain_map(self, y_hat):
          """Adds IQN network to output of network.
  
          Args:
              y_hat (torch.Tensor): Input tensor.
                  - Univariate: [B, h, 1]
                  - Multivariate: [B, h, N]
  
          Returns:
              torch.Tensor: Domain mapped tensor.
          """
          if self.eval() and self.has_predicted:
              quantiles = torch.full(
                  size=y_hat.shape,
                  fill_value=self.q,
                  device=y_hat.device,
                  dtype=y_hat.dtype,
              )
              quantiles = quantiles.unsqueeze(-1)
          else:
              quantiles = self._sample_quantiles(
                  sample_size=y_hat.shape, device=y_hat.device
              )
  
          # Embed the quantiles and add to y_hat
          emb_taus = self.quantile_layer(quantiles)
          emb_inputs = y_hat.unsqueeze(-1) * (1.0 + emb_taus)
          emb_outputs = self.output_layer(emb_inputs)
  
          # Domain map
          y_hat = emb_outputs.squeeze(-1)
  
          return y_hat
  
  
  def weighted_average(
      x: torch.Tensor, weights: Optional[torch.Tensor] = None, dim=None
  ) -> torch.Tensor:
      """Computes the weighted average of a given tensor across a given dim.
  
      Masks values associated with weight zero, meaning instead of `nan * 0 = nan`
      you will get `0 * 0 = 0`.
  
      Args:
          x (torch.Tensor): Input tensor, of which the average must be computed.
          weights (Optional[torch.Tensor], optional): Weights tensor, of the same shape as `x`. Defaults to None.
          dim (optional): The dim along which to average `x`. Defaults to None.
  
      Returns:
          torch.Tensor: The tensor with values averaged along the specified `dim`.
      """
      if weights is not None:
          weighted_tensor = torch.where(weights != 0, x * weights, torch.zeros_like(x))
          sum_weights = torch.clamp(
              weights.sum(dim=dim) if dim else weights.sum(), min=1.0
          )
          return (
              weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()
          ) / sum_weights
      else:
          return x.mean(dim=dim)
  
  
  def bernoulli_scale_decouple(output, loc=None, scale=None):
      """Bernoulli Scale Decouple.
  
      Stabilizes model's output optimization, by learning residual
      variance and residual location based on anchoring `loc`, `scale`.
      Also adds Bernoulli domain protection to the distribution parameters.
  
      Args:
          output: Model output tensor.
          loc (optional): Location parameter. Defaults to None.
          scale (optional): Scale parameter. Defaults to None.
  
      Returns:
          tuple: Processed probabilities.
      """
      probs = output[0]
      # if (loc is not None) and (scale is not None):
      #    rate = (rate * scale) + loc
      probs = F.sigmoid(probs)  # .clone()
      return (probs,)
  
  
  def student_scale_decouple(output, loc=None, scale=None, eps: float = 0.1):
      """Student-T Scale Decouple.
  
      Stabilizes model's output optimization, by learning residual
      variance and residual location based on anchoring `loc`, `scale`.
      Also adds StudentT domain protection to the distribution parameters.
  
      Args:
          output: Model output tensor.
          loc (optional): Location parameter. Defaults to None.
          scale (optional): Scale parameter. Defaults to None.
          eps (float, optional): Epsilon value for numerical stability. Defaults to 0.1.
  
      Returns:
          tuple: Processed degrees of freedom, mean, and scale parameters.
      """
      df, mean, tscale = output
      tscale = F.softplus(tscale)
      if (loc is not None) and (scale is not None):
  