From 930f85560609e4a9cab1bc306b067dad8bd66cb5 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 10 Mar 2026 00:33:48 +0500
Subject: [PATCH 01/23] [ENH] Units_v2 Model added

---
 pytorch_forecasting/layers/_units/__init__.py |  11 +
 pytorch_forecasting/layers/_units/_units.py   | 127 ++++++++++
 pytorch_forecasting/models/__init__.py        |   2 +
 pytorch_forecasting/models/units/__init__.py  |   8 +
 .../models/units/_units_pkg_v2.py             |  56 +++++
 pytorch_forecasting/models/units/_units_v2.py | 203 ++++++++++++++++
 tests/test_models/test_units_v2.py            | 225 ++++++++++++++++++
 7 files changed, 632 insertions(+)
 create mode 100644 pytorch_forecasting/layers/_units/__init__.py
 create mode 100644 pytorch_forecasting/layers/_units/_units.py
 create mode 100644 pytorch_forecasting/models/units/__init__.py
 create mode 100644 pytorch_forecasting/models/units/_units_pkg_v2.py
 create mode 100644 pytorch_forecasting/models/units/_units_v2.py
 create mode 100644 tests/test_models/test_units_v2.py

diff --git a/pytorch_forecasting/layers/_units/__init__.py b/pytorch_forecasting/layers/_units/__init__.py
new file mode 100644
index 000000000..2755ce971
--- /dev/null
+++ b/pytorch_forecasting/layers/_units/__init__.py
@@ -0,0 +1,11 @@
+"""
+UniTS layer abstractions.
+"""
+
+from pytorch_forecasting.layers._units._units import (
+    _PatchEmbedding,
+    _PositionalEncoding,
+    _TransformerBlock,
+)
+
+__all__ = ["_PatchEmbedding", "_PositionalEncoding", "_TransformerBlock"]
diff --git a/pytorch_forecasting/layers/_units/_units.py b/pytorch_forecasting/layers/_units/_units.py
new file mode 100644
index 000000000..1d4855d42
--- /dev/null
+++ b/pytorch_forecasting/layers/_units/_units.py
@@ -0,0 +1,127 @@
+"""
+Core Neural Network Layers for the UniTS architecture.
+"""
+
+import math
+
+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)
+
+
+class _PositionalEncoding(nn.Module):
+    """
+    Sinusoidal positional encoding.
+
+    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__()
+        self.drop = nn.Dropout(dropout)
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        half = d_model // 2
+        div_term = torch.exp(
+            torch.arange(0, half, dtype=torch.float) * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term[: pe[:, 0::2].size(1)])
+        pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].size(1)])
+        self.register_buffer("pe", pe.unsqueeze(0))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop(x + self.pe[:, : x.size(1), :])
+
+
+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
diff --git a/pytorch_forecasting/models/__init__.py b/pytorch_forecasting/models/__init__.py
index dc635b261..3ec0e6c65 100644
--- a/pytorch_forecasting/models/__init__.py
+++ b/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_pkg_v2
 from pytorch_forecasting.models.xlstm import xLSTMTime
 
 __all__ = [
@@ -41,5 +42,6 @@
     "DecoderMLP",
     "TiDEModel",
     "TimeXer",
+    "UniTS_pkg_v2",
     "xLSTMTime",
 ]
diff --git a/pytorch_forecasting/models/units/__init__.py b/pytorch_forecasting/models/units/__init__.py
new file mode 100644
index 000000000..30debc0ed
--- /dev/null
+++ b/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"]
diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
new file mode 100644
index 000000000..a04483234
--- /dev/null
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -0,0 +1,56 @@
+"""
+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
+    """
+
+    _tags = {
+        "info:name": "UniTS",
+        "authors": ["Muhammad-Rebaal"],
+        "capability:exogenous": True,
+        "capability:multivariate": True,
+        "capability:pred_int": False,
+        "capability:flexible_history_length": 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._tslib_data_module import TslibDataModule
+
+        return TslibDataModule
+
+    @classmethod
+    def get_test_train_params(cls):
+        """Define varied configurations for auto-testing."""
+        return [
+            {
+                "patch_len": 8,
+                "stride": 4,
+                "datamodule_cfg": {"context_length": 12, "prediction_length": 4},
+            },
+            {
+                "d_model": 32,
+                "n_heads": 4,
+                "patch_len": 8,
+                "stride": 4,
+                "datamodule_cfg": {"context_length": 12, "prediction_length": 4},
+            },
+            {
+                "patch_len": 8,
+                "stride": 4,
+                "datamodule_cfg": {"context_length": 16, "prediction_length": 4},
+            },
+        ]
diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
new file mode 100644
index 000000000..cda7017fe
--- /dev/null
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -0,0 +1,203 @@
+"""
+UniTS: Unified Time Series Model for PyTorch Forecasting.
+"""
+
+from typing import Any
+import warnings
+
+import torch
+import torch.nn as nn
+from torch.optim import Optimizer
+
+from pytorch_forecasting.layers._units import (
+    _PatchEmbedding,
+    _PositionalEncoding,
+    _TransformerBlock,
+)
+from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel
+
+
+class UniTS(TslibBaseModel):
+    """
+    UniTS: Unified Time Series Model.
+
+    Patch-based transformer for multivariate time series forecasting.
+    Implements a simplified version of the architecture from the paper, using
+    channel-independent patching: each channel is projected separately with a
+    shared linear layer, then averaged across channels.
+
+    Parameters
+    ----------
+    loss : nn.Module
+        Loss function for training.
+    d_model : int, optional
+        Transformer model dimension. Default is 64.
+    n_heads : int, optional
+        Number of self-attention heads. Must divide d_model. Default is 8.
+    e_layers : int, optional
+        Number of transformer encoder layers. Default is 3.
+    d_ff : int, optional
+        Feed-forward hidden dimension. Default is 512.
+    dropout : float, optional
+        Dropout probability. Default is 0.1.
+    patch_len : int, optional
+        Patch length in time steps. Must be <= context_length. Default is 16.
+    stride : int, optional
+        Stride between patches. Default is 8.
+    prompt_len : int, optional
+        Number of learnable task-prompt tokens prepended to patch sequence.
+        Default is 10.
+    logging_metrics : list[nn.Module] or None, optional
+        Metrics to log during training.
+    optimizer : Optimizer or str or None, optional
+        Optimizer. Default is 'adam'.
+    optimizer_params : dict or None, optional
+        Optimizer parameters.
+    lr_scheduler : str or None, optional
+        Learning rate scheduler.
+    lr_scheduler_params : dict or None, optional
+        Scheduler parameters.
+    metadata : dict or None, optional
+        Dataset metadata provided by TslibDataModule.
+    """
+
+    @classmethod
+    def _pkg(cls):
+        """Package containing the model."""
+        from pytorch_forecasting.models.units._units_pkg_v2 import UniTS_pkg_v2
+
+        return UniTS_pkg_v2
+
+    def __init__(
+        self,
+        loss: nn.Module,
+        d_model: int = 64,
+        n_heads: int = 8,
+        e_layers: int = 3,
+        d_ff: int = 512,
+        dropout: float = 0.1,
+        patch_len: int = 16,
+        stride: int = 8,
+        prompt_len: int = 10,
+        logging_metrics: list[nn.Module] | None = None,
+        optimizer: Optimizer | str | None = "adam",
+        optimizer_params: dict | None = None,
+        lr_scheduler: str | None = None,
+        lr_scheduler_params: dict | None = None,
+        metadata: dict | None = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            loss=loss,
+            logging_metrics=logging_metrics,
+            optimizer=optimizer,
+            optimizer_params=optimizer_params,
+            lr_scheduler=lr_scheduler,
+            lr_scheduler_params=lr_scheduler_params,
+            metadata=metadata,
+        )
+
+        warnings.warn(
+            "UniTS is an experimental model implemented on TslibBaseModelV2. "
+            "It is an unstable version and may be subject to unannounced changes. "
+            "Please use with caution.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+        self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"])
+
+        if d_model % n_heads != 0:
+            raise ValueError(
+                f"d_model ({d_model}) must be divisible by n_heads ({n_heads})."
+            )
+        if patch_len > self.context_length:
+            raise ValueError(
+                f"patch_len ({patch_len}) must not exceed "
+                f"context_length ({self.context_length})."
+            )
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.e_layers = e_layers
+        self.d_ff = d_ff
+        self.dropout = dropout
+        self.patch_len = patch_len
+        self.stride = stride
+        self.prompt_len = prompt_len
+
+        self._init_network()
+
+    def _init_network(self):
+        """Initialise model layers."""
+        self.num_patches = max(
+            1, (self.context_length - self.patch_len) // self.stride + 1
+        )
+
+        self.patch_embedding = _PatchEmbedding(
+            patch_len=self.patch_len,
+            stride=self.stride,
+            d_model=self.d_model,
+            dropout=self.dropout,
+        )
+
+        # Learnable Task Prompt Tokens
+        self.prompt_tokens = nn.Parameter(torch.empty(1, self.prompt_len, self.d_model))
+        nn.init.trunc_normal_(self.prompt_tokens, std=0.02)
+
+        self.pos_enc = _PositionalEncoding(
+            d_model=self.d_model,
+            max_len=self.prompt_len + self.num_patches + 16,
+            dropout=self.dropout,
+        )
+
+        self.encoder = nn.ModuleList(
+            [
+                _TransformerBlock(self.d_model, self.n_heads, self.d_ff, self.dropout)
+                for _ in range(self.e_layers)
+            ]
+        )
+
+        self.norm = nn.LayerNorm(self.d_model)
+
+        self.head = nn.Sequential(
+            nn.Flatten(start_dim=1),
+            nn.Linear(
+                self.num_patches * self.d_model,
+                self.prediction_length * self.target_dim,
+            ),
+        )
+
+    def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        """
+        Forward logic passing data through the abstracted layers.
+        """
+        target = x["history_target"]
+        B = target.size(0)
+
+        cont = x.get("history_cont")
+        if cont is not None and cont.size(-1) > 0:
+            src = torch.cat([cont, target], dim=-1)
+        else:
+            src = target
+
+        mean = src.mean(dim=1, keepdim=True)
+        std = src.std(dim=1, keepdim=True, unbiased=False) + 1e-5
+        src = (src - mean) / std
+
+        patch_emb = self.patch_embedding(src)
+
+        seq = torch.cat([self.prompt_tokens.expand(B, -1, -1), patch_emb], dim=1)
+        seq = self.pos_enc(seq)
+
+        for layer in self.encoder:
+            seq = layer(seq)
+
+        seq = self.norm(seq)
+        patch_out = seq[:, self.prompt_len : self.prompt_len + self.num_patches, :]
+        out = self.head(patch_out).view(B, self.prediction_length, self.target_dim)
+
+        target_mean = mean[:, :, -self.target_dim :]
+        target_std = std[:, :, -self.target_dim :]
+
+        return {"prediction": out * target_std + target_mean}
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
new file mode 100644
index 000000000..1094887a9
--- /dev/null
+++ b/tests/test_models/test_units_v2.py
@@ -0,0 +1,225 @@
+"""Tests for UniTS v2 model."""
+
+import inspect
+
+import pytest
+import torch
+import torch.nn as nn
+
+from pytorch_forecasting.models.units import UniTS_pkg_v2
+from pytorch_forecasting.models.units._units_v2 import UniTS
+
+
+def _make_metadata(
+    context_length=24,
+    prediction_length=6,
+    target_dim=3,
+    cont_dim=0,
+    cat_dim=0,
+    features="M",
+):
+    return {
+        "context_length": context_length,
+        "prediction_length": prediction_length,
+        "features": features,
+        "n_features": {
+            "target": target_dim,
+            "continuous": cont_dim,
+            "categorical": cat_dim,
+            "static_categorical": 0,
+            "static_continuous": 0,
+        },
+        "feature_indices": {
+            "target": list(range(target_dim)),
+            "continuous": list(range(cont_dim)),
+            "categorical": [],
+            "known": [],
+            "unknown": [],
+        },
+        "feature_names": {},
+    }
+
+
+def _make_model(metadata, **kwargs):
+    return UniTS(loss=nn.MSELoss(), metadata=metadata, **kwargs)
+
+
+def _make_batch(metadata, batch_size=2):
+    B = batch_size
+    T = metadata["context_length"]
+    C = metadata["n_features"]["target"]
+    Cc = metadata["n_features"]["continuous"]
+    batch = {
+        "history_target": torch.randn(B, T, C),
+        "history_time_idx": torch.arange(T).unsqueeze(0).expand(B, -1),
+        "target_scale": {
+            "scale": torch.ones(B, 1, C),
+            "center": torch.zeros(B, 1, C),
+        },
+    }
+    if Cc > 0:
+        batch["history_cont"] = torch.randn(B, T, Cc)
+    return batch
+
+
+class TestUniTSPkg:
+    """Tests for UniTS_pkg_v2."""
+
+    def test_get_cls(self):
+        assert UniTS_pkg_v2.get_cls().__name__ == "UniTS"
+
+    def test_pkg_tags(self):
+        instance = UniTS_pkg_v2()
+        assert instance.get_tag("info:name") == "UniTS"
+        assert instance.get_tag("capability:multivariate") is True
+        assert instance.get_tag("capability:exogenous") is True
+        assert instance.get_tag("capability:pred_int") is False
+
+    def test_get_datamodule_cls(self):
+        assert UniTS_pkg_v2.get_datamodule_cls() is not None
+
+    def test_get_test_train_params(self):
+        params = UniTS_pkg_v2.get_test_train_params()
+        assert isinstance(params, list) and len(params) > 0
+        for p in params:
+            assert "datamodule_cfg" in p
+            dm = p["datamodule_cfg"]
+            assert "context_length" in dm
+            assert "prediction_length" in dm
+            # patch_len must not exceed context_length
+            patch_len = p.get("patch_len", 16)
+            assert (
+                patch_len <= dm["context_length"]
+            ), f"patch_len={patch_len} exceeds context_length={dm['context_length']}"
+
+    def test_get_test_train_params_independent(self):
+        """Each param dict must be independent — no shared mutable objects."""
+        params = UniTS_pkg_v2.get_test_train_params()
+        dm_ids = [id(p["datamodule_cfg"]) for p in params]
+        assert len(dm_ids) == len(
+            set(dm_ids)
+        ), "datamodule_cfg dicts are shared objects"
+
+
+class TestUniTSForward:
+    """Forward pass shape and numerical correctness."""
+
+    @pytest.fixture
+    def meta(self):
+        return _make_metadata()
+
+    def test_output_shape_default(self, meta):
+        model = _make_model(meta)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert "prediction" in out
+        assert out["prediction"].shape == (2, 6, 3)
+
+    def test_output_shape_small_dmodel(self, meta):
+        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert out["prediction"].shape == (2, 6, 3)
+
+    def test_output_shape_single_layer(self, meta):
+        model = _make_model(meta, e_layers=1, d_model=32, n_heads=4, d_ff=64)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert out["prediction"].shape == (2, 6, 3)
+
+    def test_no_nan(self, meta):
+        model = _make_model(meta)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert not torch.isnan(out["prediction"]).any()
+
+    def test_no_inf(self, meta):
+        model = _make_model(meta)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert not torch.isinf(out["prediction"]).any()
+
+    def test_pkg_classmethod(self):
+        assert UniTS._pkg().__name__ == "UniTS_pkg_v2"
+
+    def test_gradients_flow(self, meta):
+        """Loss.backward() must not raise and must produce non-zero gradients."""
+        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
+        model.train()
+        out = model(_make_batch(meta))
+        target = torch.randn(2, 6, 3)
+        loss = nn.MSELoss()(out["prediction"], target)
+        loss.backward()
+        grads = [p.grad for p in model.parameters() if p.grad is not None]
+        assert len(grads) > 0
+        assert all(not torch.isnan(g).any() for g in grads)
+
+    def test_exogenous_features(self):
+        """Model must accept and use continuous exogenous features."""
+        meta = _make_metadata(
+            context_length=24, prediction_length=6, target_dim=2, cont_dim=3
+        )
+        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert out["prediction"].shape == (2, 6, 2)
+
+    @pytest.mark.parametrize("pred_len", [1, 3, 12, 24])
+    def test_prediction_lengths(self, pred_len):
+        meta = _make_metadata(
+            context_length=48, prediction_length=pred_len, target_dim=2
+        )
+        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert out["prediction"].shape == (2, pred_len, 2)
+
+    @pytest.mark.parametrize("target_dim", [1, 4, 7])
+    def test_target_dims(self, target_dim):
+        meta = _make_metadata(
+            context_length=24, prediction_length=6, target_dim=target_dim
+        )
+        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
+        model.eval()
+        with torch.no_grad():
+            out = model(_make_batch(meta))
+        assert out["prediction"].shape == (2, 6, target_dim)
+
+
+class TestUniTSParams:
+    """Parameter validation."""
+
+    def test_d_model_not_divisible_by_n_heads(self):
+        with pytest.raises(ValueError, match="d_model"):
+            UniTS(loss=nn.MSELoss(), metadata=_make_metadata(), d_model=33, n_heads=8)
+
+    def test_patch_len_exceeds_context_length(self):
+        with pytest.raises(ValueError, match="patch_len"):
+            UniTS(
+                loss=nn.MSELoss(),
+                metadata=_make_metadata(context_length=8),
+                patch_len=16,
+            )
+
+    def test_default_hyperparameters(self):
+        sig = inspect.signature(UniTS.__init__)
+        defaults = {
+            k: v.default
+            for k, v in sig.parameters.items()
+            if v.default is not inspect.Parameter.empty
+        }
+        assert defaults["d_model"] == 64
+        assert defaults["n_heads"] == 8
+        assert defaults["e_layers"] == 3
+        assert defaults["d_ff"] == 512
+        assert defaults["dropout"] == 0.1
+        assert defaults["patch_len"] == 16
+        assert defaults["stride"] == 8
+        assert defaults["prompt_len"] == 10

From 2d46f32dae0b0f7e852bc597ae7e352ed12cf070 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 16 Mar 2026 17:50:10 +0500
Subject: [PATCH 02/23] Add UniTS model, its package container, and integrate
 it into the models initialization

---
 pytorch_forecasting/models/__init__.py            | 4 ++--
 pytorch_forecasting/models/units/_units_pkg_v2.py | 2 +-
 pytorch_forecasting/models/units/_units_v2.py     | 4 ++++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/pytorch_forecasting/models/__init__.py b/pytorch_forecasting/models/__init__.py
index 3ec0e6c65..c2886abb9 100644
--- a/pytorch_forecasting/models/__init__.py
+++ b/pytorch_forecasting/models/__init__.py
@@ -20,7 +20,7 @@
 )
 from pytorch_forecasting.models.tide import TiDEModel
 from pytorch_forecasting.models.timexer import TimeXer
-from pytorch_forecasting.models.units import UniTS_pkg_v2
+from pytorch_forecasting.models.units import UniTS
 from pytorch_forecasting.models.xlstm import xLSTMTime
 
 __all__ = [
@@ -42,6 +42,6 @@
     "DecoderMLP",
     "TiDEModel",
     "TimeXer",
-    "UniTS_pkg_v2",
+    "UniTS",
     "xLSTMTime",
 ]
diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index a04483234..55961a8c0 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -13,7 +13,7 @@ class UniTS_pkg_v2(Base_pkg):
 
     _tags = {
         "info:name": "UniTS",
-        "authors": ["Muhammad-Rebaal"],
+        "authors": ["Muhammad-Rebaal", "sohamukute"],
         "capability:exogenous": True,
         "capability:multivariate": True,
         "capability:pred_int": False,
diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
index cda7017fe..644d3e9fd 100644
--- a/pytorch_forecasting/models/units/_units_v2.py
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -21,6 +21,10 @@ class UniTS(TslibBaseModel):
     """
     UniTS: Unified Time Series Model.
 
+    GitHub Repository : https://github.com/mims-harvard/UniTS
+
+    Research Paper : https://arxiv.org/abs/2403.00131
+
     Patch-based transformer for multivariate time series forecasting.
     Implements a simplified version of the architecture from the paper, using
     channel-independent patching: each channel is projected separately with a

From e9cd22f90887a663df01b01709392f4268421a7e Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 17 Mar 2026 10:48:55 +0500
Subject: [PATCH 03/23] [ENH] Added a default fixture

---
 .../models/units/_units_pkg_v2.py             | 27 ++++++++++++++++---
 1 file changed, 23 insertions(+), 4 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 55961a8c0..181476990 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -34,19 +34,28 @@ def get_datamodule_cls(cls):
 
     @classmethod
     def get_test_train_params(cls):
-        """Define varied configurations for auto-testing."""
-        return [
+        """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,
-                "datamodule_cfg": {"context_length": 12, "prediction_length": 4},
             },
             {
                 "d_model": 32,
                 "n_heads": 4,
                 "patch_len": 8,
                 "stride": 4,
-                "datamodule_cfg": {"context_length": 12, "prediction_length": 4},
             },
             {
                 "patch_len": 8,
@@ -54,3 +63,13 @@ def get_test_train_params(cls):
                 "datamodule_cfg": {"context_length": 16, "prediction_length": 4},
             },
         ]
+
+        default_dm_cfg = {"context_length": 12, "prediction_length": 4}
+
+        for param in params:
+            current_dm_cfg = param.get("datamodule_cfg", {})
+            default_dm_cfg.update(current_dm_cfg)
+
+            param["datamodule_cfg"] = default_dm_cfg
+
+        return params

From c8c81a9fa1ab25f473440e3d5098c37a0eb8241b Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 17 Mar 2026 11:49:40 +0500
Subject: [PATCH 04/23] feat: Add `UniTS_pkg_v2` for UniTS model definition,
 metadata, and test parameters.

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 181476990..2ea7955d7 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -64,12 +64,11 @@ def get_test_train_params(cls):
             },
         ]
 
-        default_dm_cfg = {"context_length": 12, "prediction_length": 4}
+        base_dm_cfg = {"context_length": 12, "prediction_length": 4}
 
         for param in params:
-            current_dm_cfg = param.get("datamodule_cfg", {})
-            default_dm_cfg.update(current_dm_cfg)
-
-            param["datamodule_cfg"] = default_dm_cfg
+            merged = base_dm_cfg.copy()
+            merged.update(param.get("datamodule_cfg", {}))
+            param["datamodule_cfg"] = merged
 
         return params

From 61e4d294065fcff9d69be17ad2b65384662ccc3b Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 17 Mar 2026 12:14:39 +0500
Subject: [PATCH 05/23] [BUG] Fix default fixture context_length and shared
 dict

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 2ea7955d7..3149fc9d4 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -64,7 +64,7 @@ def get_test_train_params(cls):
             },
         ]
 
-        base_dm_cfg = {"context_length": 12, "prediction_length": 4}
+        base_dm_cfg = {"context_length": 16, "prediction_length": 4}
 
         for param in params:
             merged = base_dm_cfg.copy()

From 8900902712a64724fb0da6af800fc1ed2a1b16e5 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 23 Mar 2026 20:16:26 +0500
Subject: [PATCH 06/23] fix: Code Refactored

---
 pytorch_forecasting/layers/__init__.py        |   7 +-
 .../layers/_blocks/__init__.py                |   3 +-
 .../layers/_blocks/_transformer_block.py      |  45 +++++++
 .../layers/_embeddings/__init__.py            |   4 +
 .../layers/_embeddings/_patch_embedding.py    |  55 ++++++++
 .../_embeddings/_positional_embedding.py      |  26 ++++
 pytorch_forecasting/layers/_units/__init__.py |  10 +-
 pytorch_forecasting/layers/_units/_units.py   | 127 ------------------
 pytorch_forecasting/models/units/_units_v2.py |   9 +-
 9 files changed, 146 insertions(+), 140 deletions(-)
 create mode 100644 pytorch_forecasting/layers/_blocks/_transformer_block.py
 create mode 100644 pytorch_forecasting/layers/_embeddings/_patch_embedding.py
 delete mode 100644 pytorch_forecasting/layers/_units/_units.py

diff --git a/pytorch_forecasting/layers/__init__.py b/pytorch_forecasting/layers/__init__.py
index ffaaec653..61a32299e 100644
--- a/pytorch_forecasting/layers/__init__.py
+++ b/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",
diff --git a/pytorch_forecasting/layers/_blocks/__init__.py b/pytorch_forecasting/layers/_blocks/__init__.py
index 512760a31..8bd1114ac 100644
--- a/pytorch_forecasting/layers/_blocks/__init__.py
+++ b/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"]
diff --git a/pytorch_forecasting/layers/_blocks/_transformer_block.py b/pytorch_forecasting/layers/_blocks/_transformer_block.py
new file mode 100644
index 000000000..be64292aa
--- /dev/null
+++ b/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
diff --git a/pytorch_forecasting/layers/_embeddings/__init__.py b/pytorch_forecasting/layers/_embeddings/__init__.py
index 7e1977bc9..e680df6de 100644
--- a/pytorch_forecasting/layers/_embeddings/__init__.py
+++ b/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",
diff --git a/pytorch_forecasting/layers/_embeddings/_patch_embedding.py b/pytorch_forecasting/layers/_embeddings/_patch_embedding.py
new file mode 100644
index 000000000..aa56deb4d
--- /dev/null
+++ b/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)
diff --git a/pytorch_forecasting/layers/_embeddings/_positional_embedding.py b/pytorch_forecasting/layers/_embeddings/_positional_embedding.py
index 70b8dbe74..3611c8b47 100644
--- a/pytorch_forecasting/layers/_embeddings/_positional_embedding.py
+++ b/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), :])
diff --git a/pytorch_forecasting/layers/_units/__init__.py b/pytorch_forecasting/layers/_units/__init__.py
index 2755ce971..73b8afc72 100644
--- a/pytorch_forecasting/layers/_units/__init__.py
+++ b/pytorch_forecasting/layers/_units/__init__.py
@@ -1,11 +1,11 @@
 """
-UniTS layer abstractions.
+UniTS layer abstractions - re-exported from canonical locations.
 """
 
-from pytorch_forecasting.layers._units._units import (
-    _PatchEmbedding,
-    _PositionalEncoding,
-    _TransformerBlock,
+from pytorch_forecasting.layers._blocks._transformer_block import _TransformerBlock
+from pytorch_forecasting.layers._embeddings._patch_embedding import _PatchEmbedding
+from pytorch_forecasting.layers._embeddings._positional_embedding import (
+    _PositionalEmbedding as _PositionalEncoding,  # backward compat alias
 )
 
 __all__ = ["_PatchEmbedding", "_PositionalEncoding", "_TransformerBlock"]
diff --git a/pytorch_forecasting/layers/_units/_units.py b/pytorch_forecasting/layers/_units/_units.py
deleted file mode 100644
index 1d4855d42..000000000
--- a/pytorch_forecasting/layers/_units/_units.py
+++ /dev/null
@@ -1,127 +0,0 @@
-"""
-Core Neural Network Layers for the UniTS architecture.
-"""
-
-import math
-
-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)
-
-
-class _PositionalEncoding(nn.Module):
-    """
-    Sinusoidal positional encoding.
-
-    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__()
-        self.drop = nn.Dropout(dropout)
-        pe = torch.zeros(max_len, d_model)
-        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
-        half = d_model // 2
-        div_term = torch.exp(
-            torch.arange(0, half, dtype=torch.float) * (-math.log(10000.0) / d_model)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term[: pe[:, 0::2].size(1)])
-        pe[:, 1::2] = torch.cos(position * div_term[: pe[:, 1::2].size(1)])
-        self.register_buffer("pe", pe.unsqueeze(0))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.drop(x + self.pe[:, : x.size(1), :])
-
-
-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
diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
index 644d3e9fd..8bfa22574 100644
--- a/pytorch_forecasting/models/units/_units_v2.py
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -9,11 +9,8 @@
 import torch.nn as nn
 from torch.optim import Optimizer
 
-from pytorch_forecasting.layers._units import (
-    _PatchEmbedding,
-    _PositionalEncoding,
-    _TransformerBlock,
-)
+from pytorch_forecasting.layers._blocks import _TransformerBlock
+from pytorch_forecasting.layers._embeddings import _PatchEmbedding, _PositionalEmbedding
 from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel
 
 
@@ -149,7 +146,7 @@ def _init_network(self):
         self.prompt_tokens = nn.Parameter(torch.empty(1, self.prompt_len, self.d_model))
         nn.init.trunc_normal_(self.prompt_tokens, std=0.02)
 
-        self.pos_enc = _PositionalEncoding(
+        self.pos_enc = _PositionalEmbedding(
             d_model=self.d_model,
             max_len=self.prompt_len + self.num_patches + 16,
             dropout=self.dropout,

From a7c0172b55b740f20f78c1dbd64eef725cc5799c Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 16 Jun 2026 04:45:43 +0500
Subject: [PATCH 07/23] feat: add UniTS_pkg_v2 remaining tags and update the
 import

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 3149fc9d4..ef616d64c 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -9,15 +9,21 @@ 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",
-        "authors": ["Muhammad-Rebaal", "sohamukute"],
+        "info:pred_type": ["point"],
+        "info:y_type": ["numeric"],
+        "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
@@ -28,7 +34,7 @@ def get_cls(cls):
 
     @classmethod
     def get_datamodule_cls(cls):
-        from pytorch_forecasting.data._tslib_data_module import TslibDataModule
+        from pytorch_forecasting.data.data_module import TslibDataModule
 
         return TslibDataModule
 

From 08f46b04b0ac3f338527752858712ea8aa757442 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 16 Jun 2026 05:27:19 +0500
Subject: [PATCH 08/23] fix : units_v2 test error

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index ef616d64c..c0e3f8cef 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -39,7 +39,7 @@ def get_datamodule_cls(cls):
         return TslibDataModule
 
     @classmethod
-    def get_test_train_params(cls):
+    def get_base_test_params(cls):
         """Return testing parameter settings for the trainer.
 
         Returns

From 49c27f870a2707a3f4c9ca29454a4736960ad915 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 16 Jun 2026 11:54:46 +0500
Subject: [PATCH 09/23] fix: pytest error

---
 tests/test_models/test_units_v2.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 1094887a9..832a7be8e 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -78,8 +78,8 @@ def test_pkg_tags(self):
     def test_get_datamodule_cls(self):
         assert UniTS_pkg_v2.get_datamodule_cls() is not None
 
-    def test_get_test_train_params(self):
-        params = UniTS_pkg_v2.get_test_train_params()
+    def get_base_test_params(self):
+        params = UniTS_pkg_v2.get_base_test_params()
         assert isinstance(params, list) and len(params) > 0
         for p in params:
             assert "datamodule_cfg" in p
@@ -94,7 +94,7 @@ def test_get_test_train_params(self):
 
     def test_get_test_train_params_independent(self):
         """Each param dict must be independent — no shared mutable objects."""
-        params = UniTS_pkg_v2.get_test_train_params()
+        params = UniTS_pkg_v2.get_base_test_params()
         dm_ids = [id(p["datamodule_cfg"]) for p in params]
         assert len(dm_ids) == len(
             set(dm_ids)

From 349c72691f2910233a3e14ce10f4dfec32659b5a Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Tue, 16 Jun 2026 12:10:02 +0500
Subject: [PATCH 10/23] fix: pytest error

---
 tests/test_models/test_units_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 832a7be8e..30a6bfb2c 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -78,7 +78,7 @@ def test_pkg_tags(self):
     def test_get_datamodule_cls(self):
         assert UniTS_pkg_v2.get_datamodule_cls() is not None
 
-    def get_base_test_params(self):
+    def test_get_base_test_params(self):
         params = UniTS_pkg_v2.get_base_test_params()
         assert isinstance(params, list) and len(params) > 0
         for p in params:

From 858a1464587e57345712f527951fb72d35e3a33d Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 22 Jun 2026 17:28:34 +0500
Subject: [PATCH 11/23] fix:pytest hallucination

---
 .../models/units/_units_pkg_v2.py             |   2 +-
 tests/test_models/test_units_v2.py            | 378 ++++++++----------
 2 files changed, 165 insertions(+), 215 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index c0e3f8cef..ef616d64c 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -39,7 +39,7 @@ def get_datamodule_cls(cls):
         return TslibDataModule
 
     @classmethod
-    def get_base_test_params(cls):
+    def get_test_train_params(cls):
         """Return testing parameter settings for the trainer.
 
         Returns
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 30a6bfb2c..3cc7e178d 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -1,225 +1,175 @@
 """Tests for UniTS v2 model."""
 
-import inspect
-
+import numpy as np
+import pandas as pd
 import pytest
 import torch
-import torch.nn as nn
 
-from pytorch_forecasting.models.units import UniTS_pkg_v2
+from pytorch_forecasting.data import TimeSeries
+from pytorch_forecasting.data.data_module import TslibDataModule
+from pytorch_forecasting.metrics import MAE, SMAPE
 from pytorch_forecasting.models.units._units_v2 import UniTS
 
 
-def _make_metadata(
-    context_length=24,
-    prediction_length=6,
-    target_dim=3,
-    cont_dim=0,
-    cat_dim=0,
-    features="M",
-):
-    return {
-        "context_length": context_length,
-        "prediction_length": prediction_length,
-        "features": features,
-        "n_features": {
-            "target": target_dim,
-            "continuous": cont_dim,
-            "categorical": cat_dim,
-            "static_categorical": 0,
-            "static_continuous": 0,
-        },
-        "feature_indices": {
-            "target": list(range(target_dim)),
-            "continuous": list(range(cont_dim)),
-            "categorical": [],
-            "known": [],
-            "unknown": [],
-        },
-        "feature_names": {},
-    }
-
-
-def _make_model(metadata, **kwargs):
-    return UniTS(loss=nn.MSELoss(), metadata=metadata, **kwargs)
-
-
-def _make_batch(metadata, batch_size=2):
-    B = batch_size
-    T = metadata["context_length"]
-    C = metadata["n_features"]["target"]
-    Cc = metadata["n_features"]["continuous"]
-    batch = {
-        "history_target": torch.randn(B, T, C),
-        "history_time_idx": torch.arange(T).unsqueeze(0).expand(B, -1),
-        "target_scale": {
-            "scale": torch.ones(B, 1, C),
-            "center": torch.zeros(B, 1, C),
-        },
-    }
-    if Cc > 0:
-        batch["history_cont"] = torch.randn(B, T, Cc)
-    return batch
-
-
-class TestUniTSPkg:
-    """Tests for UniTS_pkg_v2."""
-
-    def test_get_cls(self):
-        assert UniTS_pkg_v2.get_cls().__name__ == "UniTS"
-
-    def test_pkg_tags(self):
-        instance = UniTS_pkg_v2()
-        assert instance.get_tag("info:name") == "UniTS"
-        assert instance.get_tag("capability:multivariate") is True
-        assert instance.get_tag("capability:exogenous") is True
-        assert instance.get_tag("capability:pred_int") is False
-
-    def test_get_datamodule_cls(self):
-        assert UniTS_pkg_v2.get_datamodule_cls() is not None
-
-    def test_get_base_test_params(self):
-        params = UniTS_pkg_v2.get_base_test_params()
-        assert isinstance(params, list) and len(params) > 0
-        for p in params:
-            assert "datamodule_cfg" in p
-            dm = p["datamodule_cfg"]
-            assert "context_length" in dm
-            assert "prediction_length" in dm
-            # patch_len must not exceed context_length
-            patch_len = p.get("patch_len", 16)
-            assert (
-                patch_len <= dm["context_length"]
-            ), f"patch_len={patch_len} exceeds context_length={dm['context_length']}"
-
-    def test_get_test_train_params_independent(self):
-        """Each param dict must be independent — no shared mutable objects."""
-        params = UniTS_pkg_v2.get_base_test_params()
-        dm_ids = [id(p["datamodule_cfg"]) for p in params]
-        assert len(dm_ids) == len(
-            set(dm_ids)
-        ), "datamodule_cfg dicts are shared objects"
-
-
-class TestUniTSForward:
-    """Forward pass shape and numerical correctness."""
-
-    @pytest.fixture
-    def meta(self):
-        return _make_metadata()
-
-    def test_output_shape_default(self, meta):
-        model = _make_model(meta)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert "prediction" in out
-        assert out["prediction"].shape == (2, 6, 3)
-
-    def test_output_shape_small_dmodel(self, meta):
-        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert out["prediction"].shape == (2, 6, 3)
-
-    def test_output_shape_single_layer(self, meta):
-        model = _make_model(meta, e_layers=1, d_model=32, n_heads=4, d_ff=64)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert out["prediction"].shape == (2, 6, 3)
-
-    def test_no_nan(self, meta):
-        model = _make_model(meta)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert not torch.isnan(out["prediction"]).any()
-
-    def test_no_inf(self, meta):
-        model = _make_model(meta)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert not torch.isinf(out["prediction"]).any()
-
-    def test_pkg_classmethod(self):
-        assert UniTS._pkg().__name__ == "UniTS_pkg_v2"
-
-    def test_gradients_flow(self, meta):
-        """Loss.backward() must not raise and must produce non-zero gradients."""
-        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
-        model.train()
-        out = model(_make_batch(meta))
-        target = torch.randn(2, 6, 3)
-        loss = nn.MSELoss()(out["prediction"], target)
-        loss.backward()
-        grads = [p.grad for p in model.parameters() if p.grad is not None]
-        assert len(grads) > 0
-        assert all(not torch.isnan(g).any() for g in grads)
-
-    def test_exogenous_features(self):
-        """Model must accept and use continuous exogenous features."""
-        meta = _make_metadata(
-            context_length=24, prediction_length=6, target_dim=2, cont_dim=3
+@pytest.fixture
+def sample_multivariate_data():
+    """Sample multivariate data for testing."""
+    np.random.seed(42)
+    series_len = 30
+    num_groups = 3
+    data = []
+
+    for i in range(num_groups):
+        time_idx = np.arange(series_len, dtype=np.int64)
+        trend = 100 + i * 20 + 0.5 * time_idx
+        seasonal = 10 * np.sin(2 * np.pi * time_idx / 12)
+        noise = np.random.normal(0, 5, series_len)
+        target = trend + seasonal + noise
+
+        temperature = (
+            20
+            + 15 * np.sin(2 * np.pi * time_idx / 365)
+            + np.random.normal(0, 3, series_len)
         )
-        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert out["prediction"].shape == (2, 6, 2)
-
-    @pytest.mark.parametrize("pred_len", [1, 3, 12, 24])
-    def test_prediction_lengths(self, pred_len):
-        meta = _make_metadata(
-            context_length=48, prediction_length=pred_len, target_dim=2
+        humidity = (
+            30
+            + 20 * np.cos(2 * np.pi * time_idx / 7)
+            + np.random.normal(0, 5, series_len)
         )
-        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert out["prediction"].shape == (2, pred_len, 2)
-
-    @pytest.mark.parametrize("target_dim", [1, 4, 7])
-    def test_target_dims(self, target_dim):
-        meta = _make_metadata(
-            context_length=24, prediction_length=6, target_dim=target_dim
+
+        df_group = pd.DataFrame(
+            {
+                "time_idx": time_idx,
+                "group_id": f"group_{i}",
+                "value": target.astype(np.float32),
+                "temperature": temperature.astype(np.float32),
+                "humidity": humidity.astype(np.float32),
+            }
         )
-        model = _make_model(meta, d_model=32, n_heads=4, d_ff=64)
-        model.eval()
-        with torch.no_grad():
-            out = model(_make_batch(meta))
-        assert out["prediction"].shape == (2, 6, target_dim)
-
-
-class TestUniTSParams:
-    """Parameter validation."""
-
-    def test_d_model_not_divisible_by_n_heads(self):
-        with pytest.raises(ValueError, match="d_model"):
-            UniTS(loss=nn.MSELoss(), metadata=_make_metadata(), d_model=33, n_heads=8)
-
-    def test_patch_len_exceeds_context_length(self):
-        with pytest.raises(ValueError, match="patch_len"):
-            UniTS(
-                loss=nn.MSELoss(),
-                metadata=_make_metadata(context_length=8),
-                patch_len=16,
-            )
-
-    def test_default_hyperparameters(self):
-        sig = inspect.signature(UniTS.__init__)
-        defaults = {
-            k: v.default
-            for k, v in sig.parameters.items()
-            if v.default is not inspect.Parameter.empty
-        }
-        assert defaults["d_model"] == 64
-        assert defaults["n_heads"] == 8
-        assert defaults["e_layers"] == 3
-        assert defaults["d_ff"] == 512
-        assert defaults["dropout"] == 0.1
-        assert defaults["patch_len"] == 16
-        assert defaults["stride"] == 8
-        assert defaults["prompt_len"] == 10
+        data.append(df_group)
+
+    df = pd.concat(data, ignore_index=True)
+    df["group_id"] = df["group_id"].astype("category")
+    return df
+
+
+@pytest.fixture
+def basic_timeseries_dataset(sample_multivariate_data):
+    """Create a basic TimeSeries dataset for testing."""
+    return TimeSeries(
+        data=sample_multivariate_data,
+        time="time_idx",
+        target="value",
+        group=["group_id"],
+        num=["value", "temperature", "humidity"],
+        cat=[],
+        known=["temperature", "humidity", "time_idx"],
+        static=[],
+    )
+
+
+@pytest.fixture
+def basic_tslib_data_module(basic_timeseries_dataset):
+    """Create a basic TslibDataModule for testing."""
+    return TslibDataModule(
+        time_series_dataset=basic_timeseries_dataset,
+        batch_size=2,
+        context_length=16,
+        prediction_length=4,
+        train_val_test_split=(0.7, 0.15, 0.15),
+    )
+
+
+@pytest.fixture
+def basic_metadata(basic_tslib_data_module):
+    """Basic metadata from data module for model initialization."""
+    basic_tslib_data_module.setup()
+    return basic_tslib_data_module.metadata
+
+
+@pytest.fixture(params=[16, 32], ids=["d_model_16", "d_model_32"])
+def model(request, basic_metadata):
+    """Initialize a UniTS model for testing."""
+    return UniTS(
+        loss=MAE(),
+        d_model=request.param,
+        n_heads=4,
+        e_layers=2,
+        d_ff=64,
+        dropout=0.1,
+        patch_len=8,
+        stride=4,
+        logging_metrics=[SMAPE()],
+        optimizer="adam",
+        metadata=basic_metadata,
+    )
+
+
+def test_basic_model_initialization(model, basic_metadata):
+    """Test the basic model initialization."""
+    assert isinstance(model, UniTS)
+    assert model.d_model in [16, 32]
+    assert model.n_heads == 4
+    assert model.e_layers == 2
+    assert model.patch_len == 8
+
+    assert model.context_length == basic_metadata["context_length"]
+    assert model.prediction_length == basic_metadata["prediction_length"]
+
+
+def test_multivariate_single_series(model, basic_tslib_data_module):
+    """Test forward pass shape and no NaN outputs."""
+    basic_tslib_data_module.setup()
+    train_dataloader = basic_tslib_data_module.train_dataloader()
+    batch = next(iter(train_dataloader))[0]
+
+    model.eval()
+    with torch.no_grad():
+        output = model(batch)
+
+    assert "prediction" in output
+    predictions = output["prediction"]
+
+    batch_size = batch["history_target"].shape[0]
+    assert predictions.shape == (batch_size, model.prediction_length, model.target_dim)
+    assert not torch.isnan(predictions).any()
+    assert not torch.isinf(predictions).any()
+
+
+def test_integration_with_datamodule(model, basic_tslib_data_module):
+    """Test integration of UniTS model with TslibDataModule."""
+    basic_tslib_data_module.setup(stage="fit")
+    basic_tslib_data_module.setup(stage="test")
+
+    train_loader = basic_tslib_data_module.train_dataloader()
+    val_loader = basic_tslib_data_module.val_dataloader()
+    test_loader = basic_tslib_data_module.test_dataloader()
+
+    model.eval()
+    with torch.no_grad():
+        train_batch = next(iter(train_loader))[0]
+        train_output = model(train_batch)
+        assert train_output["prediction"].shape[1] == model.prediction_length
+
+        try:
+            val_batch = next(iter(val_loader))[0]
+            val_output = model(val_batch)
+            assert val_output["prediction"].shape[1] == model.prediction_length
+        except StopIteration:
+            pass
+
+        try:
+            test_batch = next(iter(test_loader))[0]
+            test_output = model(test_batch)
+            assert test_output["prediction"].shape[1] == model.prediction_length
+        except StopIteration:
+            pass
+
+
+def test_parameter_validation(basic_metadata):
+    """Test parameter validation for UniTS."""
+    with pytest.raises(ValueError, match="d_model"):
+        UniTS(loss=MAE(), metadata=basic_metadata, d_model=33, n_heads=8)
+
+    with pytest.raises(ValueError, match="patch_len"):
+        UniTS(loss=MAE(), metadata=basic_metadata, patch_len=32)

From 1fa26b314957524c5997b6822a1743aee17d2d87 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 29 Jun 2026 14:45:57 +0500
Subject: [PATCH 12/23] fix: removed the v1 tags

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index ef616d64c..0067cf04d 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -14,8 +14,6 @@ class UniTS_pkg_v2(Base_pkg):
 
     _tags = {
         "info:name": "UniTS",
-        "info:pred_type": ["point"],
-        "info:y_type": ["numeric"],
         "info:compute": 4,
         "authors": ["Muhammad-Rebaal", "gasvn", "sohamukute"],
         "python_dependencies": ["torch"],

From 8e7919f1a7ddff869e6cdbf7e4ce6bcd908bcd25 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Fri, 3 Jul 2026 16:38:26 +0500
Subject: [PATCH 13/23] fix: Performed De-duplication

---
 pytorch_forecasting/layers/_units/__init__.py | 11 ----
 tests/test_models/test_units_v2.py            | 61 -------------------
 2 files changed, 72 deletions(-)
 delete mode 100644 pytorch_forecasting/layers/_units/__init__.py

diff --git a/pytorch_forecasting/layers/_units/__init__.py b/pytorch_forecasting/layers/_units/__init__.py
deleted file mode 100644
index 73b8afc72..000000000
--- a/pytorch_forecasting/layers/_units/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-"""
-UniTS layer abstractions - re-exported from canonical locations.
-"""
-
-from pytorch_forecasting.layers._blocks._transformer_block import _TransformerBlock
-from pytorch_forecasting.layers._embeddings._patch_embedding import _PatchEmbedding
-from pytorch_forecasting.layers._embeddings._positional_embedding import (
-    _PositionalEmbedding as _PositionalEncoding,  # backward compat alias
-)
-
-__all__ = ["_PatchEmbedding", "_PositionalEncoding", "_TransformerBlock"]
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 3cc7e178d..1898f0f51 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -105,67 +105,6 @@ def model(request, basic_metadata):
     )
 
 
-def test_basic_model_initialization(model, basic_metadata):
-    """Test the basic model initialization."""
-    assert isinstance(model, UniTS)
-    assert model.d_model in [16, 32]
-    assert model.n_heads == 4
-    assert model.e_layers == 2
-    assert model.patch_len == 8
-
-    assert model.context_length == basic_metadata["context_length"]
-    assert model.prediction_length == basic_metadata["prediction_length"]
-
-
-def test_multivariate_single_series(model, basic_tslib_data_module):
-    """Test forward pass shape and no NaN outputs."""
-    basic_tslib_data_module.setup()
-    train_dataloader = basic_tslib_data_module.train_dataloader()
-    batch = next(iter(train_dataloader))[0]
-
-    model.eval()
-    with torch.no_grad():
-        output = model(batch)
-
-    assert "prediction" in output
-    predictions = output["prediction"]
-
-    batch_size = batch["history_target"].shape[0]
-    assert predictions.shape == (batch_size, model.prediction_length, model.target_dim)
-    assert not torch.isnan(predictions).any()
-    assert not torch.isinf(predictions).any()
-
-
-def test_integration_with_datamodule(model, basic_tslib_data_module):
-    """Test integration of UniTS model with TslibDataModule."""
-    basic_tslib_data_module.setup(stage="fit")
-    basic_tslib_data_module.setup(stage="test")
-
-    train_loader = basic_tslib_data_module.train_dataloader()
-    val_loader = basic_tslib_data_module.val_dataloader()
-    test_loader = basic_tslib_data_module.test_dataloader()
-
-    model.eval()
-    with torch.no_grad():
-        train_batch = next(iter(train_loader))[0]
-        train_output = model(train_batch)
-        assert train_output["prediction"].shape[1] == model.prediction_length
-
-        try:
-            val_batch = next(iter(val_loader))[0]
-            val_output = model(val_batch)
-            assert val_output["prediction"].shape[1] == model.prediction_length
-        except StopIteration:
-            pass
-
-        try:
-            test_batch = next(iter(test_loader))[0]
-            test_output = model(test_batch)
-            assert test_output["prediction"].shape[1] == model.prediction_length
-        except StopIteration:
-            pass
-
-
 def test_parameter_validation(basic_metadata):
     """Test parameter validation for UniTS."""
     with pytest.raises(ValueError, match="d_model"):

From 7eb7fcee4d4a839cde448cf361a5976171bf184a Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Fri, 3 Jul 2026 16:44:25 +0500
Subject: [PATCH 14/23] Added API reference for the Units model

---
 docs/source/m_layer_v2.rst | 1 +
 docs/source/pkg_v2.rst     | 1 +
 2 files changed, 2 insertions(+)

diff --git a/docs/source/m_layer_v2.rst b/docs/source/m_layer_v2.rst
index 00f4f0ac4..e21a324fb 100644
--- a/docs/source/m_layer_v2.rst
+++ b/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
diff --git a/docs/source/pkg_v2.rst b/docs/source/pkg_v2.rst
index 560ad7a3c..a4b937120 100644
--- a/docs/source/pkg_v2.rst
+++ b/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

From 0dc7a8dc9fdb3e5a1d831f40de083e5fefc12a4e Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 6 Jul 2026 22:46:54 +0500
Subject: [PATCH 15/23] feat: Converted the model on the BaseClass instead of
 the TslibModel Class

---
 .../models/units/_units_pkg_v2.py              | 13 +++++++++----
 pytorch_forecasting/models/units/_units_v2.py  | 18 +++++++++++-------
 tests/test_models/test_units_v2.py             | 18 +++++++++---------
 3 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 0067cf04d..f190fe617 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -32,9 +32,11 @@ def get_cls(cls):
 
     @classmethod
     def get_datamodule_cls(cls):
-        from pytorch_forecasting.data.data_module import TslibDataModule
+        from pytorch_forecasting.data.data_module import (
+            EncoderDecoderTimeSeriesDataModule,
+        )
 
-        return TslibDataModule
+        return EncoderDecoderTimeSeriesDataModule
 
     @classmethod
     def get_test_train_params(cls):
@@ -64,11 +66,14 @@ def get_test_train_params(cls):
             {
                 "patch_len": 8,
                 "stride": 4,
-                "datamodule_cfg": {"context_length": 16, "prediction_length": 4},
+                "datamodule_cfg": {
+                    "max_encoder_length": 16,
+                    "max_prediction_length": 4,
+                },
             },
         ]
 
-        base_dm_cfg = {"context_length": 16, "prediction_length": 4}
+        base_dm_cfg = {"max_encoder_length": 16, "max_prediction_length": 4}
 
         for param in params:
             merged = base_dm_cfg.copy()
diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
index 8bfa22574..1ae6f42f1 100644
--- a/pytorch_forecasting/models/units/_units_v2.py
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -11,10 +11,10 @@
 
 from pytorch_forecasting.layers._blocks import _TransformerBlock
 from pytorch_forecasting.layers._embeddings import _PatchEmbedding, _PositionalEmbedding
-from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel
+from pytorch_forecasting.models.base._base_model_v2 import BaseModel
 
 
-class UniTS(TslibBaseModel):
+class UniTS(BaseModel):
     """
     UniTS: Unified Time Series Model.
 
@@ -59,7 +59,7 @@ class UniTS(TslibBaseModel):
     lr_scheduler_params : dict or None, optional
         Scheduler parameters.
     metadata : dict or None, optional
-        Dataset metadata provided by TslibDataModule.
+        Dataset metadata provided by EncoderDecoderTimeSeriesDataModule.
     """
 
     @classmethod
@@ -95,11 +95,10 @@ def __init__(
             optimizer_params=optimizer_params,
             lr_scheduler=lr_scheduler,
             lr_scheduler_params=lr_scheduler_params,
-            metadata=metadata,
         )
 
         warnings.warn(
-            "UniTS is an experimental model implemented on TslibBaseModelV2. "
+            "UniTS is an experimental model implemented on BaseModel. "
             "It is an unstable version and may be subject to unannounced changes. "
             "Please use with caution.",
             UserWarning,
@@ -108,6 +107,11 @@ def __init__(
 
         self.save_hyperparameters(ignore=["loss", "logging_metrics", "metadata"])
 
+        self.metadata = metadata or {}
+        self.context_length = self.metadata.get("max_encoder_length", 0)
+        self.prediction_length = self.metadata.get("max_prediction_length", 0)
+        self.target_dim = self.metadata.get("target", 1)
+
         if d_model % n_heads != 0:
             raise ValueError(
                 f"d_model ({d_model}) must be divisible by n_heads ({n_heads})."
@@ -173,10 +177,10 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         """
         Forward logic passing data through the abstracted layers.
         """
-        target = x["history_target"]
+        target = x["target_past"]
         B = target.size(0)
 
-        cont = x.get("history_cont")
+        cont = x.get("encoder_cont")
         if cont is not None and cont.size(-1) > 0:
             src = torch.cat([cont, target], dim=-1)
         else:
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 1898f0f51..378af53c5 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -6,7 +6,7 @@
 import torch
 
 from pytorch_forecasting.data import TimeSeries
-from pytorch_forecasting.data.data_module import TslibDataModule
+from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
 from pytorch_forecasting.metrics import MAE, SMAPE
 from pytorch_forecasting.models.units._units_v2 import UniTS
 
@@ -69,22 +69,22 @@ def basic_timeseries_dataset(sample_multivariate_data):
 
 
 @pytest.fixture
-def basic_tslib_data_module(basic_timeseries_dataset):
-    """Create a basic TslibDataModule for testing."""
-    return TslibDataModule(
+def basic_data_module(basic_timeseries_dataset):
+    """Create a basic DataModule for testing."""
+    return EncoderDecoderTimeSeriesDataModule(
         time_series_dataset=basic_timeseries_dataset,
         batch_size=2,
-        context_length=16,
-        prediction_length=4,
+        max_encoder_length=16,
+        max_prediction_length=4,
         train_val_test_split=(0.7, 0.15, 0.15),
     )
 
 
 @pytest.fixture
-def basic_metadata(basic_tslib_data_module):
+def basic_metadata(basic_data_module):
     """Basic metadata from data module for model initialization."""
-    basic_tslib_data_module.setup()
-    return basic_tslib_data_module.metadata
+    basic_data_module.setup()
+    return basic_data_module.metadata
 
 
 @pytest.fixture(params=[16, 32], ids=["d_model_16", "d_model_32"])

From 96c229ef62c14ec83ae2d88fbb73722b3602e271 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 27 Jul 2026 19:43:12 +0500
Subject: [PATCH 16/23] remove the non-important tag, losses suport added, and
 tests added

---
 .../models/units/_units_pkg_v2.py             |  19 +-
 pytorch_forecasting/models/units/_units_v2.py |  50 +--
 tests/test_models/test_units_v2.py            | 293 ++++++++++++++++--
 3 files changed, 312 insertions(+), 50 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index f190fe617..a97f1d168 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -16,11 +16,12 @@ class UniTS_pkg_v2(Base_pkg):
         "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,
+        "info:pred_type": ["point", "quantile", "distribution"],
+        "info:y_type": ["numeric"],
+        "capability:pred_int": True,
+        "capability:flexible_history_length": True,
         "capability:cold_start": False,
     }
 
@@ -51,6 +52,8 @@ def get_test_train_params(cls):
             instance. ``create_test_instance`` uses the first (or only) dictionary in
             ``params``.
         """
+        from pytorch_forecasting.metrics import NormalDistributionLoss, QuantileLoss
+
         params = [
             {},
             {
@@ -71,6 +74,16 @@ def get_test_train_params(cls):
                     "max_prediction_length": 4,
                 },
             },
+            {
+                "patch_len": 8,
+                "stride": 4,
+                "loss": QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
+            },
+            {
+                "patch_len": 8,
+                "stride": 4,
+                "loss": NormalDistributionLoss(),
+            },
         ]
 
         base_dm_cfg = {"max_encoder_length": 16, "max_prediction_length": 4}
diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
index 1ae6f42f1..ccb62b108 100644
--- a/pytorch_forecasting/models/units/_units_v2.py
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -11,6 +11,7 @@
 
 from pytorch_forecasting.layers._blocks import _TransformerBlock
 from pytorch_forecasting.layers._embeddings import _PatchEmbedding, _PositionalEmbedding
+from pytorch_forecasting.metrics import DistributionLoss, QuantileLoss
 from pytorch_forecasting.models.base._base_model_v2 import BaseModel
 
 
@@ -135,9 +136,6 @@ def __init__(
 
     def _init_network(self):
         """Initialise model layers."""
-        self.num_patches = max(
-            1, (self.context_length - self.patch_len) // self.stride + 1
-        )
 
         self.patch_embedding = _PatchEmbedding(
             patch_len=self.patch_len,
@@ -150,12 +148,6 @@ def _init_network(self):
         self.prompt_tokens = nn.Parameter(torch.empty(1, self.prompt_len, self.d_model))
         nn.init.trunc_normal_(self.prompt_tokens, std=0.02)
 
-        self.pos_enc = _PositionalEmbedding(
-            d_model=self.d_model,
-            max_len=self.prompt_len + self.num_patches + 16,
-            dropout=self.dropout,
-        )
-
         self.encoder = nn.ModuleList(
             [
                 _TransformerBlock(self.d_model, self.n_heads, self.d_ff, self.dropout)
@@ -163,14 +155,27 @@ def _init_network(self):
             ]
         )
 
+        self.pos_enc = _PositionalEmbedding(d_model=self.d_model, dropout=self.dropout)
         self.norm = nn.LayerNorm(self.d_model)
 
+        self.n_quantiles = None
+        self.n_dist_args = None
+
+        if isinstance(self.loss, QuantileLoss):
+            self.n_quantiles = len(self.loss.quantiles)
+
+        elif isinstance(self.loss, DistributionLoss):
+            self.n_dist_args = len(self.loss.distribution_arguments)
+        output_dim = self.prediction_length * self.target_dim
+
+        if self.n_quantiles is not None:
+            output_dim = self.prediction_length * self.target_dim * self.n_quantiles
+        elif self.n_dist_args is not None:
+            output_dim = self.prediction_length * self.target_dim * self.n_dist_args
+
         self.head = nn.Sequential(
             nn.Flatten(start_dim=1),
-            nn.Linear(
-                self.num_patches * self.d_model,
-                self.prediction_length * self.target_dim,
-            ),
+            nn.Linear(self.prompt_len * self.d_model, output_dim),
         )
 
     def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
@@ -186,10 +191,6 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         else:
             src = target
 
-        mean = src.mean(dim=1, keepdim=True)
-        std = src.std(dim=1, keepdim=True, unbiased=False) + 1e-5
-        src = (src - mean) / std
-
         patch_emb = self.patch_embedding(src)
 
         seq = torch.cat([self.prompt_tokens.expand(B, -1, -1), patch_emb], dim=1)
@@ -199,10 +200,15 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
             seq = layer(seq)
 
         seq = self.norm(seq)
-        patch_out = seq[:, self.prompt_len : self.prompt_len + self.num_patches, :]
-        out = self.head(patch_out).view(B, self.prediction_length, self.target_dim)
+        patch_out = seq[:, : self.prompt_len, :]
+
+        raw = self.head(patch_out)
 
-        target_mean = mean[:, :, -self.target_dim :]
-        target_std = std[:, :, -self.target_dim :]
+        if self.n_quantiles is not None:
+            out = raw.view(B, self.prediction_length, self.target_dim, self.n_quantiles)
+        elif self.n_dist_args is not None:
+            out = raw.view(B, self.prediction_length, self.target_dim, self.n_dist_args)
+        else:
+            out = raw.view(B, self.prediction_length, self.target_dim)
 
-        return {"prediction": out * target_std + target_mean}
+        return {"prediction": out}
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 378af53c5..e2032ec57 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -7,13 +7,23 @@
 
 from pytorch_forecasting.data import TimeSeries
 from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
-from pytorch_forecasting.metrics import MAE, SMAPE
+from pytorch_forecasting.metrics import MAE, SMAPE, NormalDistributionLoss, QuantileLoss
 from pytorch_forecasting.models.units._units_v2 import UniTS
 
+BATCH_SIZE = 2
+MAX_ENCODER_LENGTH = 16
+MAX_PREDICTION_LENGTH = 4
+D_MODEL = 16
+N_HEADS = 4
+E_LAYERS = 1
+D_FF = 32
+PATCH_LEN = 8
+STRIDE = 4
+
 
 @pytest.fixture
 def sample_multivariate_data():
-    """Sample multivariate data for testing."""
+    """Synthetic multivariate time series DataFrame."""
     np.random.seed(42)
     series_len = 30
     num_groups = 3
@@ -55,7 +65,7 @@ def sample_multivariate_data():
 
 @pytest.fixture
 def basic_timeseries_dataset(sample_multivariate_data):
-    """Create a basic TimeSeries dataset for testing."""
+    """TimeSeries object from sample data."""
     return TimeSeries(
         data=sample_multivariate_data,
         time="time_idx",
@@ -70,45 +80,278 @@ def basic_timeseries_dataset(sample_multivariate_data):
 
 @pytest.fixture
 def basic_data_module(basic_timeseries_dataset):
-    """Create a basic DataModule for testing."""
+    """EncoderDecoderTimeSeriesDataModule, not yet set up."""
     return EncoderDecoderTimeSeriesDataModule(
         time_series_dataset=basic_timeseries_dataset,
-        batch_size=2,
-        max_encoder_length=16,
-        max_prediction_length=4,
+        batch_size=BATCH_SIZE,
+        max_encoder_length=MAX_ENCODER_LENGTH,
+        max_prediction_length=MAX_PREDICTION_LENGTH,
         train_val_test_split=(0.7, 0.15, 0.15),
     )
 
 
 @pytest.fixture
 def basic_metadata(basic_data_module):
-    """Basic metadata from data module for model initialization."""
+    """Metadata dict extracted after DataModule setup."""
     basic_data_module.setup()
     return basic_data_module.metadata
 
 
-@pytest.fixture(params=[16, 32], ids=["d_model_16", "d_model_32"])
-def model(request, basic_metadata):
-    """Initialize a UniTS model for testing."""
-    return UniTS(
+def test_basic_attributes(basic_metadata):
+    """Model attributes match constructor args."""
+    model = UniTS(
         loss=MAE(),
-        d_model=request.param,
-        n_heads=4,
-        e_layers=2,
-        d_ff=64,
-        dropout=0.1,
-        patch_len=8,
-        stride=4,
-        logging_metrics=[SMAPE()],
-        optimizer="adam",
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        e_layers=E_LAYERS,
+        d_ff=D_FF,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
         metadata=basic_metadata,
     )
+    assert model.d_model == D_MODEL
+    assert model.n_heads == N_HEADS
+    assert model.e_layers == E_LAYERS
+    assert model.context_length == MAX_ENCODER_LENGTH
+    assert model.prediction_length == MAX_PREDICTION_LENGTH
 
 
-def test_parameter_validation(basic_metadata):
-    """Test parameter validation for UniTS."""
+def test_d_model_not_divisible_by_n_heads(basic_metadata):
+    """d_model % n_heads != 0 must raise ValueError."""
     with pytest.raises(ValueError, match="d_model"):
-        UniTS(loss=MAE(), metadata=basic_metadata, d_model=33, n_heads=8)
+        UniTS(
+            loss=MAE(),
+            d_model=33,
+            n_heads=8,
+            metadata=basic_metadata,
+        )
+
 
+def test_patch_len_exceeds_context(basic_metadata):
+    """patch_len > context_length must raise ValueError."""
     with pytest.raises(ValueError, match="patch_len"):
-        UniTS(loss=MAE(), metadata=basic_metadata, patch_len=32)
+        UniTS(
+            loss=MAE(),
+            patch_len=MAX_ENCODER_LENGTH + 1,
+            metadata=basic_metadata,
+        )
+
+
+def test_hyperparameters_saved(basic_metadata):
+    """save_hyperparameters stores model config (not loss/metadata)."""
+    model = UniTS(
+        loss=MAE(),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    assert model.hparams["d_model"] == D_MODEL
+    assert "loss" not in model.hparams
+    assert "metadata" not in model.hparams
+
+
+def test_output_shape_point_loss(basic_metadata, basic_data_module):
+    """Prediction shape is (B, pred_len, target_dim) with point loss."""
+    basic_data_module.setup()
+    model = UniTS(
+        loss=MAE(),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        e_layers=E_LAYERS,
+        d_ff=D_FF,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, _ = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        output = model(batch_x)
+
+    pred = output["prediction"]
+    actual_batch = batch_x["target_past"].shape[0]
+    assert pred.shape == (
+        actual_batch,
+        MAX_PREDICTION_LENGTH,
+        basic_metadata["target"],
+    )
+
+
+def test_no_nan_or_inf(basic_metadata, basic_data_module):
+    """Output must not contain NaN or Inf values."""
+    basic_data_module.setup()
+    model = UniTS(
+        loss=MAE(),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, _ = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        pred = model(batch_x)["prediction"]
+
+    assert not torch.isnan(pred).any(), "Predictions contain NaN"
+    assert not torch.isinf(pred).any(), "Predictions contain Inf"
+
+
+def test_quantile_loss_output_shape(basic_metadata, basic_data_module):
+    """QuantileLoss must produce (B, pred_len, target_dim, n_quantiles) output."""
+    basic_data_module.setup()
+    quantiles = [0.1, 0.5, 0.9]
+    model = UniTS(
+        loss=QuantileLoss(quantiles=quantiles),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, _ = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        pred = model(batch_x)["prediction"]
+
+    actual_batch = batch_x["target_past"].shape[0]
+    assert pred.shape == (
+        actual_batch,
+        MAX_PREDICTION_LENGTH,
+        basic_metadata["target"],
+        len(quantiles),
+    )
+
+
+def test_quantile_n_quantiles_attribute(basic_metadata):
+    """n_quantiles attribute set correctly when using QuantileLoss."""
+    model = UniTS(
+        loss=QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    assert model.n_quantiles == 3
+
+
+def test_point_loss_n_quantiles_is_none(basic_metadata):
+    """n_quantiles is None when using a point loss like MAE."""
+    model = UniTS(
+        loss=MAE(),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    assert model.n_quantiles is None
+
+
+def test_distribution_loss_output_shape(basic_metadata, basic_data_module):
+    """DistributionLoss must produce (B, pred_len, target_dim, n_dist_args) output."""
+    basic_data_module.setup()
+    loss = NormalDistributionLoss()
+    model = UniTS(
+        loss=loss,
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, _ = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        pred = model(batch_x)["prediction"]
+
+    actual_batch = batch_x["target_past"].shape[0]
+    assert pred.shape == (
+        actual_batch,
+        MAX_PREDICTION_LENGTH,
+        basic_metadata["target"],
+        len(loss.distribution_arguments),
+    )
+
+
+@pytest.mark.parametrize("loss_cls", [MAE, SMAPE])
+def test_multiple_point_losses(loss_cls, basic_metadata, basic_data_module):
+    """Model produces valid output with various point losses."""
+    basic_data_module.setup()
+    model = UniTS(
+        loss=loss_cls(),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, _ = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        pred = model(batch_x)["prediction"]
+
+    assert not torch.isnan(pred).any()
+
+
+def test_metadata_dimensions(basic_metadata):
+    """Metadata contains all keys the model constructor reads."""
+    required_keys = [
+        "max_encoder_length",
+        "max_prediction_length",
+        "target",
+    ]
+    for key in required_keys:
+        assert key in basic_metadata, f"Missing metadata key: {key}"
+
+
+def test_train_batch_roundtrip(basic_metadata, basic_data_module):
+    """Model processes a real training batch and returns valid output."""
+    basic_data_module.setup()
+    model = UniTS(
+        loss=MAE(),
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, batch_y = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        output = model(batch_x)
+
+    pred = output["prediction"]
+    assert pred.shape[1] == MAX_PREDICTION_LENGTH
+    assert not torch.isnan(pred).any()
+
+
+def test_loss_computes_on_real_batch(basic_metadata, basic_data_module):
+    """Loss function returns a finite scalar on a real batch."""
+    basic_data_module.setup()
+    loss_fn = MAE()
+    model = UniTS(
+        loss=loss_fn,
+        d_model=D_MODEL,
+        n_heads=N_HEADS,
+        patch_len=PATCH_LEN,
+        stride=STRIDE,
+        metadata=basic_metadata,
+    )
+    model.eval()
+
+    batch_x, batch_y = next(iter(basic_data_module.train_dataloader()))
+    with torch.no_grad():
+        pred = model(batch_x)["prediction"]
+        loss_val = loss_fn(pred, batch_y)
+
+    assert torch.isfinite(loss_val), "Loss is not finite"

From 36f2e6a381ee57a9b8147238a78b7dc737ca55bf Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 27 Jul 2026 22:51:43 +0500
Subject: [PATCH 17/23] fix pytest error

---
 .../models/units/_units_pkg_v2.py             |  2 +-
 pytorch_forecasting/models/units/_units_v2.py | 11 ++++++++--
 tests/test_models/test_units_v2.py            | 22 +++++++++----------
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index a97f1d168..1d9f89dfe 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -40,7 +40,7 @@ def get_datamodule_cls(cls):
         return EncoderDecoderTimeSeriesDataModule
 
     @classmethod
-    def get_test_train_params(cls):
+    def get_base_test_params(cls):
         """Return testing parameter settings for the trainer.
 
         Returns
diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
index ccb62b108..162b8ed21 100644
--- a/pytorch_forecasting/models/units/_units_v2.py
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -205,10 +205,17 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         raw = self.head(patch_out)
 
         if self.n_quantiles is not None:
-            out = raw.view(B, self.prediction_length, self.target_dim, self.n_quantiles)
+            if self.target_dim == 1:
+                out = raw.view(B, self.prediction_length, self.n_quantiles)
+            else:
+                out = raw.view(B, self.prediction_length, self.target_dim, self.n_quantiles)
         elif self.n_dist_args is not None:
-            out = raw.view(B, self.prediction_length, self.target_dim, self.n_dist_args)
+            if self.target_dim == 1:
+                out = raw.view(B, self.prediction_length, self.n_dist_args)
+            else:
+                out = raw.view(B, self.prediction_length, self.target_dim, self.n_dist_args)
         else:
             out = raw.view(B, self.prediction_length, self.target_dim)
 
+
         return {"prediction": out}
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index e2032ec57..3e7c092f8 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -220,12 +220,12 @@ def test_quantile_loss_output_shape(basic_metadata, basic_data_module):
         pred = model(batch_x)["prediction"]
 
     actual_batch = batch_x["target_past"].shape[0]
-    assert pred.shape == (
-        actual_batch,
-        MAX_PREDICTION_LENGTH,
-        basic_metadata["target"],
-        len(quantiles),
+    expected_shape = (
+        (actual_batch, MAX_PREDICTION_LENGTH, len(quantiles))
+        if basic_metadata["target"] == 1
+        else (actual_batch, MAX_PREDICTION_LENGTH, basic_metadata["target"], len(quantiles))
     )
+    assert pred.shape == expected_shape
 
 
 def test_quantile_n_quantiles_attribute(basic_metadata):
@@ -255,7 +255,7 @@ def test_point_loss_n_quantiles_is_none(basic_metadata):
 
 
 def test_distribution_loss_output_shape(basic_metadata, basic_data_module):
-    """DistributionLoss must produce (B, pred_len, target_dim, n_dist_args) output."""
+    """DistributionLoss must produce (B, pred_len, [target_dim], n_dist_args) output."""
     basic_data_module.setup()
     loss = NormalDistributionLoss()
     model = UniTS(
@@ -273,12 +273,12 @@ def test_distribution_loss_output_shape(basic_metadata, basic_data_module):
         pred = model(batch_x)["prediction"]
 
     actual_batch = batch_x["target_past"].shape[0]
-    assert pred.shape == (
-        actual_batch,
-        MAX_PREDICTION_LENGTH,
-        basic_metadata["target"],
-        len(loss.distribution_arguments),
+    expected_shape = (
+        (actual_batch, MAX_PREDICTION_LENGTH, len(loss.distribution_arguments))
+        if basic_metadata["target"] == 1
+        else (actual_batch, MAX_PREDICTION_LENGTH, basic_metadata["target"], len(loss.distribution_arguments))
     )
+    assert pred.shape == expected_shape
 
 
 @pytest.mark.parametrize("loss_cls", [MAE, SMAPE])

From dccfe8cd7f2575c4d3e897442d6da79c78657734 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 27 Jul 2026 23:39:30 +0500
Subject: [PATCH 18/23] fix code quality

---
 pytorch_forecasting/models/units/_units_v2.py |  9 ++++++---
 tests/test_models/test_units_v2.py            | 14 ++++++++++++--
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_v2.py b/pytorch_forecasting/models/units/_units_v2.py
index 162b8ed21..5ff9a1211 100644
--- a/pytorch_forecasting/models/units/_units_v2.py
+++ b/pytorch_forecasting/models/units/_units_v2.py
@@ -208,14 +208,17 @@ def forward(self, x: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
             if self.target_dim == 1:
                 out = raw.view(B, self.prediction_length, self.n_quantiles)
             else:
-                out = raw.view(B, self.prediction_length, self.target_dim, self.n_quantiles)
+                out = raw.view(
+                    B, self.prediction_length, self.target_dim, self.n_quantiles
+                )
         elif self.n_dist_args is not None:
             if self.target_dim == 1:
                 out = raw.view(B, self.prediction_length, self.n_dist_args)
             else:
-                out = raw.view(B, self.prediction_length, self.target_dim, self.n_dist_args)
+                out = raw.view(
+                    B, self.prediction_length, self.target_dim, self.n_dist_args
+                )
         else:
             out = raw.view(B, self.prediction_length, self.target_dim)
 
-
         return {"prediction": out}
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 3e7c092f8..39128c2b9 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -223,7 +223,12 @@ def test_quantile_loss_output_shape(basic_metadata, basic_data_module):
     expected_shape = (
         (actual_batch, MAX_PREDICTION_LENGTH, len(quantiles))
         if basic_metadata["target"] == 1
-        else (actual_batch, MAX_PREDICTION_LENGTH, basic_metadata["target"], len(quantiles))
+        else (
+            actual_batch,
+            MAX_PREDICTION_LENGTH,
+            basic_metadata["target"],
+            len(quantiles),
+        )
     )
     assert pred.shape == expected_shape
 
@@ -276,7 +281,12 @@ def test_distribution_loss_output_shape(basic_metadata, basic_data_module):
     expected_shape = (
         (actual_batch, MAX_PREDICTION_LENGTH, len(loss.distribution_arguments))
         if basic_metadata["target"] == 1
-        else (actual_batch, MAX_PREDICTION_LENGTH, basic_metadata["target"], len(loss.distribution_arguments))
+        else (
+            actual_batch,
+            MAX_PREDICTION_LENGTH,
+            basic_metadata["target"],
+            len(loss.distribution_arguments),
+        )
     )
     assert pred.shape == expected_shape
 

From 3c3b66226ae263bf51aedc0a2ba8a2f1867aeb7b Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Thu, 30 Jul 2026 21:05:46 +0500
Subject: [PATCH 19/23] fix: tag removed and renamed param name

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 1d9f89dfe..ba705d29c 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -18,8 +18,6 @@ class UniTS_pkg_v2(Base_pkg):
         "authors": ["Muhammad-Rebaal", "gasvn", "sohamukute"],
         "capability:exogenous": True,
         "capability:multivariate": True,
-        "info:pred_type": ["point", "quantile", "distribution"],
-        "info:y_type": ["numeric"],
         "capability:pred_int": True,
         "capability:flexible_history_length": True,
         "capability:cold_start": False,
@@ -40,7 +38,7 @@ def get_datamodule_cls(cls):
         return EncoderDecoderTimeSeriesDataModule
 
     @classmethod
-    def get_base_test_params(cls):
+    def get_test_train_params(cls):
         """Return testing parameter settings for the trainer.
 
         Returns

From e069a2993a16cdbd74117b0da379eedcd60ff264 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Fri, 31 Jul 2026 19:08:16 +0500
Subject: [PATCH 20/23] remove dist loss as there isn't any implementation in
 base_model_v2

---
 .../models/units/_units_pkg_v2.py             |  5 ---
 tests/test_models/test_units_v2.py            | 34 +------------------
 2 files changed, 1 insertion(+), 38 deletions(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index ba705d29c..1eae61a98 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -77,11 +77,6 @@ def get_test_train_params(cls):
                 "stride": 4,
                 "loss": QuantileLoss(quantiles=[0.1, 0.5, 0.9]),
             },
-            {
-                "patch_len": 8,
-                "stride": 4,
-                "loss": NormalDistributionLoss(),
-            },
         ]
 
         base_dm_cfg = {"max_encoder_length": 16, "max_prediction_length": 4}
diff --git a/tests/test_models/test_units_v2.py b/tests/test_models/test_units_v2.py
index 39128c2b9..d41b6b4bb 100644
--- a/tests/test_models/test_units_v2.py
+++ b/tests/test_models/test_units_v2.py
@@ -7,7 +7,7 @@
 
 from pytorch_forecasting.data import TimeSeries
 from pytorch_forecasting.data.data_module import EncoderDecoderTimeSeriesDataModule
-from pytorch_forecasting.metrics import MAE, SMAPE, NormalDistributionLoss, QuantileLoss
+from pytorch_forecasting.metrics import MAE, SMAPE, QuantileLoss
 from pytorch_forecasting.models.units._units_v2 import UniTS
 
 BATCH_SIZE = 2
@@ -259,38 +259,6 @@ def test_point_loss_n_quantiles_is_none(basic_metadata):
     assert model.n_quantiles is None
 
 
-def test_distribution_loss_output_shape(basic_metadata, basic_data_module):
-    """DistributionLoss must produce (B, pred_len, [target_dim], n_dist_args) output."""
-    basic_data_module.setup()
-    loss = NormalDistributionLoss()
-    model = UniTS(
-        loss=loss,
-        d_model=D_MODEL,
-        n_heads=N_HEADS,
-        patch_len=PATCH_LEN,
-        stride=STRIDE,
-        metadata=basic_metadata,
-    )
-    model.eval()
-
-    batch_x, _ = next(iter(basic_data_module.train_dataloader()))
-    with torch.no_grad():
-        pred = model(batch_x)["prediction"]
-
-    actual_batch = batch_x["target_past"].shape[0]
-    expected_shape = (
-        (actual_batch, MAX_PREDICTION_LENGTH, len(loss.distribution_arguments))
-        if basic_metadata["target"] == 1
-        else (
-            actual_batch,
-            MAX_PREDICTION_LENGTH,
-            basic_metadata["target"],
-            len(loss.distribution_arguments),
-        )
-    )
-    assert pred.shape == expected_shape
-
-
 @pytest.mark.parametrize("loss_cls", [MAE, SMAPE])
 def test_multiple_point_losses(loss_cls, basic_metadata, basic_data_module):
     """Model produces valid output with various point losses."""

From 0de97fcdcb40d335f934609e0678e53906ee29f1 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 3 Aug 2026 17:18:49 +0500
Subject: [PATCH 21/23] fix : Code Quality

---
 docs/source/m_layer_v2.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/m_layer_v2.rst b/docs/source/m_layer_v2.rst
index d6671769c..7364234ea 100644
--- a/docs/source/m_layer_v2.rst
+++ b/docs/source/m_layer_v2.rst
@@ -48,4 +48,4 @@ See the detailed API documentation for the V2 base classes and specific model im
    models.tide._tide_dsipts._tide_v2.TIDE
    models.timexer._timexer_v2.TimeXer
    models.mlp._decodermlp_v2.DecoderMLP_v2
-   models.units._units_v2.UniTS
\ No newline at end of file
+   models.units._units_v2.UniTS

From 038d8c46481d45768ca39714236cc06b1e2f1efe Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Mon, 3 Aug 2026 17:23:48 +0500
Subject: [PATCH 22/23] fix Code Quality

---
 docs/source/pkg_v2.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/pkg_v2.rst b/docs/source/pkg_v2.rst
index 9f31316eb..24cc3841f 100644
--- a/docs/source/pkg_v2.rst
+++ b/docs/source/pkg_v2.rst
@@ -100,4 +100,4 @@ See the detailed API documentation for the available V2 Package classes below:
    models.tide._tide_dsipts._tide_v2_pkg.TIDE_pkg_v2
    models.timexer._timexer_pkg_v2.TimeXer_pkg_v2
    models.mlp._decodermlp_pkg_v2.DecoderMLP_pkg_v2
-   models.units._units_pkg_v2.UniTS_pkg_v2
\ No newline at end of file
+   models.units._units_pkg_v2.UniTS_pkg_v2

From 30b0aab16f6f59a46d8fb6155198e04998a01af3 Mon Sep 17 00:00:00 2001
From: Muhammad-Rebaal <mrebaal14@gmail.com>
Date: Sat, 8 Aug 2026 14:14:39 +0500
Subject: [PATCH 23/23] fix : Remove dead import

---
 pytorch_forecasting/models/units/_units_pkg_v2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pytorch_forecasting/models/units/_units_pkg_v2.py b/pytorch_forecasting/models/units/_units_pkg_v2.py
index 1eae61a98..29b0d4b46 100644
--- a/pytorch_forecasting/models/units/_units_pkg_v2.py
+++ b/pytorch_forecasting/models/units/_units_pkg_v2.py
@@ -50,7 +50,7 @@ def get_test_train_params(cls):
             instance. ``create_test_instance`` uses the first (or only) dictionary in
             ``params``.
         """
-        from pytorch_forecasting.metrics import NormalDistributionLoss, QuantileLoss
+        from pytorch_forecasting.metrics import QuantileLoss
 
         params = [
             {},