[ENH] Units_v2 Model added

docs/source/m_layer_v2.rst

@@ -47,3 +47,4 @@ See the detailed API documentation for the V2 base classes and specific model im
    models.samformer._samformer_v2.Samformer
    models.tide._tide_dsipts._tide_v2.TIDE
    models.timexer._timexer_v2.TimeXer
+   models.units._units_v2.UniTS

docs/source/pkg_v2.rst

@@ -99,3 +99,4 @@ See the detailed API documentation for the available V2 Package classes below:
    models.samformer._samformer_v2_pkg.Samformer_pkg_v2
    models.tide._tide_dsipts._tide_v2_pkg.TIDE_pkg_v2
    models.timexer._timexer_pkg_v2.TimeXer_pkg_v2
+   models.units._units_pkg_v2.UniTS_pkg_v2

pytorch_forecasting/layers/__init__.py

@@ -7,12 +7,14 @@
     FullAttention,
     TriangularCausalMask,
 )
-from pytorch_forecasting.layers._blocks import ResidualBlock
+from pytorch_forecasting.layers._blocks import ResidualBlock, _TransformerBlock
 from pytorch_forecasting.layers._decomposition import SeriesDecomposition
 from pytorch_forecasting.layers._embeddings import (
     DataEmbedding_inverted,
     EnEmbedding,
     PositionalEmbedding,
+    _PatchEmbedding,
+    _PositionalEmbedding,
     embedding_cat_variables,
 )
 from pytorch_forecasting.layers._encoders import (
@@ -41,6 +43,9 @@
     "DataEmbedding_inverted",
     "EnEmbedding",
     "PositionalEmbedding",
+    "_PatchEmbedding",
+    "_PositionalEmbedding",
+    "_TransformerBlock",
     "Encoder",
     "EncoderLayer",
     "FlattenHead",

pytorch_forecasting/layers/_blocks/__init__.py

@@ -1,3 +1,4 @@
 from pytorch_forecasting.layers._blocks._residual_block_dsipts import ResidualBlock
+from pytorch_forecasting.layers._blocks._transformer_block import _TransformerBlock
 
-__all__ = ["ResidualBlock"]
+__all__ = ["ResidualBlock", "_TransformerBlock"]

pytorch_forecasting/layers/_blocks/_transformer_block.py

@@ -0,0 +1,45 @@
+"""
+Pre-norm Transformer Encoder Block for PTF.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class _TransformerBlock(nn.Module):
+    """
+    Pre-norm transformer encoder block (MHSA + FFN).
+
+    Parameters
+    ----------
+    d_model : int
+        Model dimension.
+    n_heads : int
+        Number of attention heads.
+    d_ff : int
+        Feed-forward hidden dimension.
+    dropout : float
+        Dropout probability.
+    """
+
+    def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.attn = nn.MultiheadAttention(
+            embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
+        )
+        self.ff = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        normed = self.norm1(x)
+        attn_out, _ = self.attn(normed, normed, normed)
+        x = x + attn_out
+        x = x + self.ff(self.norm2(x))
+        return x

pytorch_forecasting/layers/_embeddings/__init__.py

@@ -6,13 +6,17 @@
     DataEmbedding_inverted,
 )
 from pytorch_forecasting.layers._embeddings._en_embedding import EnEmbedding
+from pytorch_forecasting.layers._embeddings._patch_embedding import _PatchEmbedding
 from pytorch_forecasting.layers._embeddings._positional_embedding import (
     PositionalEmbedding,
+    _PositionalEmbedding,
 )
 from pytorch_forecasting.layers._embeddings._sub_nn import embedding_cat_variables
 
 __all__ = [
     "PositionalEmbedding",
+    "_PositionalEmbedding",
+    "_PatchEmbedding",
     "DataEmbedding_inverted",
     "EnEmbedding",
     "embedding_cat_variables",

pytorch_forecasting/layers/_embeddings/_patch_embedding.py

@@ -0,0 +1,55 @@
+"""
+Patch Embedding Layer for PTF.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class _PatchEmbedding(nn.Module):
+    """
+    Project strided patches of a multivariate time series into d_model space.
+
+    Uses channel-independent patching: each channel's patches are projected
+    separately with a shared Linear(patch_len, d_model), then averaged across
+    channels to match the UniTS paper's channel-independent approach.
+
+    Parameters
+    ----------
+    patch_len : int
+        Length of each patch window in time steps.
+    stride : int
+        Stride between consecutive patches.
+    d_model : int
+        Output embedding dimension.
+    dropout : float
+        Dropout probability.
+    """
+
+    def __init__(self, patch_len: int, stride: int, d_model: int, dropout: float = 0.1):
+        super().__init__()
+        self.patch_len = patch_len
+        self.stride = stride
+        self.projection = nn.Linear(patch_len, d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        x : torch.Tensor
+            Shape (batch, seq_len, n_channels).
+
+        Returns
+        -------
+        torch.Tensor
+            Shape (batch, num_patches, d_model).
+        """
+        patches = x.unfold(dimension=1, size=self.patch_len, step=self.stride)
+        B, num_patches, C, P = patches.shape
+        patches = patches.permute(0, 2, 1, 3).contiguous().view(B * C, num_patches, P)
+        emb = self.drop(self.projection(patches))
+        emb = emb.view(B, C, num_patches, self.projection.out_features)
+
+        # Channel independence: average across channels as per UniTS logic
+        return emb.mean(dim=1)

pytorch_forecasting/layers/_embeddings/_positional_embedding.py

@@ -42,3 +42,29 @@ def __init__(self, d_model, max_len=5000):
 
     def forward(self, x):
         return self.pe[:, : x.size(1)]
+
+
+class _PositionalEmbedding(PositionalEmbedding):
+    """
+    Sinusoidal positional embedding with additive application and dropout.
+
+    Inherits the sinusoidal buffer from ``PositionalEmbedding`` and adds:
+    - Additive application (x + pe) in ``forward``
+    - Dropout after addition
+
+    Parameters
+    ----------
+    d_model : int
+        Embedding dimension.
+    max_len : int
+        Maximum sequence length.
+    dropout : float
+        Dropout probability.
+    """
+
+    def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
+        super().__init__(d_model, max_len)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop(x + self.pe[:, : x.size(1), :])

pytorch_forecasting/models/__init__.py

@@ -20,6 +20,7 @@
 )
 from pytorch_forecasting.models.tide import TiDEModel
 from pytorch_forecasting.models.timexer import TimeXer
+from pytorch_forecasting.models.units import UniTS
 from pytorch_forecasting.models.xlstm import xLSTMTime
 
 __all__ = [
@@ -41,5 +42,6 @@
     "DecoderMLP",
     "TiDEModel",
     "TimeXer",
+    "UniTS",
     "xLSTMTime",
 ]

pytorch_forecasting/models/units/__init__.py

@@ -0,0 +1,8 @@
+"""
+UniTS: Unified Time Series Model for time series forecasting.
+"""
+
+from pytorch_forecasting.models.units._units_pkg_v2 import UniTS_pkg_v2
+from pytorch_forecasting.models.units._units_v2 import UniTS
+
+__all__ = ["UniTS", "UniTS_pkg_v2"]

pytorch_forecasting/models/units/_units_pkg_v2.py

@@ -0,0 +1,83 @@
+"""
+Packages container for UniTS model.
+"""
+
+from pytorch_forecasting.base._base_pkg import Base_pkg
+
+
+class UniTS_pkg_v2(Base_pkg):
+    """
+    UniTS: Unified Time Series Model.
+    Reference: https://arxiv.org/abs/2403.00131
+    Github: https://github.com/mims-harvard/UniTS
+    """
+
+    _tags = {
+        "info:name": "UniTS",
+        "info:compute": 4,
+        "authors": ["Muhammad-Rebaal", "gasvn", "sohamukute"],
+        "python_dependencies": ["torch"],
+        "capability:exogenous": True,
+        "capability:multivariate": True,
+        "capability:pred_int": False,
+        "capability:flexible_history_length": False,
+        "capability:cold_start": False,
+    }
+
+    @classmethod
+    def get_cls(cls):
+        from pytorch_forecasting.models.units._units_v2 import UniTS
+
+        return UniTS
+
+    @classmethod
+    def get_datamodule_cls(cls):
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
+
+        return EncoderDecoderTimeSeriesDataModule
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Return testing parameter settings for the trainer.
+
+        Returns
+        -------
+        params : dict or list of dict, default = {}
+            Parameters to create testing instances of the class.
+            Each dict are parameters to construct an "interesting" test instance, i.e.,
+            ``MyClass(**params)`` or ``MyClass(**params[i])`` creates a valid test
+            instance. ``create_test_instance`` uses the first (or only) dictionary in
+            ``params``.
+        """
+        params = [
+            {},
+            {
+                "patch_len": 8,
+                "stride": 4,
+            },
+            {
+                "d_model": 32,
+                "n_heads": 4,
+                "patch_len": 8,
+                "stride": 4,
+            },
+            {
+                "patch_len": 8,
+                "stride": 4,
+                "datamodule_cfg": {
+                    "max_encoder_length": 16,
+                    "max_prediction_length": 4,
+                },
+            },
+        ]
+
+        base_dm_cfg = {"max_encoder_length": 16, "max_prediction_length": 4}
+
+        for param in params:
+            merged = base_dm_cfg.copy()
+            merged.update(param.get("datamodule_cfg", {}))
+            param["datamodule_cfg"] = merged
+
+        return params