[ENH] `Units_v2` Model added by Muhammad-Rebaal · Pull Request #2165 · sktime/pytorch-forecasting

Muhammad-Rebaal · 2026-03-09T19:37:46Z

Fixes #2158

Hi @fkiraly , @phoeenniixx, @PranavBhatP !

I have implemented the UniTS (Unified Time Series Model) within the PyTorch Forecasting v2 architecture.

Could you please review the PR?

Here is a summary of the changes made:

Isolated Layer Abstraction (pytorch_forecasting/layers/_units/): Implemented the core neural network components (_PatchEmbedding, _PositionalEncoding, _TransformerBlock) completely isolated from the standard base estimating logic to strictly adhere to the project's v2 layer standards.
Model Construction (pytorch_forecasting/models/units/): Created the main UniTS wrapper which inherently inherits from TslibBaseModel and handles feature input alignments, channel-independence, and forward pass routing to the abstracted layers.
V2 Package Management: Established the UniTS_pkg_v2 class representing model metadata (authors, tags) and defined automated testing configurations (get_test_train_params).
Comprehensive Test Suite: Included a full Pytest harness for the model (tests/test_models/test_units_v2.py) validating parameter exceptions, tensor shapes, gradient flow bounds, and default hyperparameter states.

Thank you!

codecov · 2026-03-09T21:36:53Z

Codecov Report

❌ Patch coverage is 95.17241% with 7 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@7bcf66c). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
pytorch_forecasting/models/units/_units_v2.py	91.13%	7 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2165   +/-   ##
=======================================
  Coverage        ?   87.57%           
=======================================
  Files           ?      180           
  Lines           ?    10308           
  Branches        ?        0           
=======================================
  Hits            ?     9027           
  Misses          ?     1281           
  Partials        ?        0

Flag	Coverage Δ
cpu	`87.57% <95.17%> (?)`
pytest	`87.57% <95.17%> (?)`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

phoeenniixx · 2026-03-15T07:19:57Z

+
+    _tags = {
+        "info:name": "UniTS",
+        "authors": ["Muhammad-Rebaal"],


I think we should also credit the actual author of the model? The one who implemented it originally?

yeah sure, I'll add him as well

phoeenniixx · 2026-03-15T07:21:32Z

 )
 from pytorch_forecasting.models.tide import TiDEModel
 from pytorch_forecasting.models.timexer import TimeXer
+from pytorch_forecasting.models.units import UniTS_pkg_v2


I think you should add the model (and not the pkg) here?

phoeenniixx · 2026-03-15T07:22:13Z

+    """
+    UniTS: Unified Time Series Model.
+
+    Patch-based transformer for multivariate time series forecasting.


I think it would be good if you could also add the references (like GH link and the paper) here?

…els initialization

Muhammad-Rebaal · 2026-03-16T12:51:23Z

Hi @phoeenniixx ,

Kindly review the PR, I've updated the PR with the requested changes.

Thank You.

…st parameters.

phoeenniixx · 2026-03-22T20:13:36Z

+import torch.nn as nn
+
+
+class _PatchEmbedding(nn.Module):


would it make sense to add it to _embeddings?

Yes I think it would make complete sense if we'd place it there in a file called _patch_embedding.py, as the _embeddings directory is already the designated home for embedding abstractions (like _data_embedding.py and _en_embedding.py).

phoeenniixx · 2026-03-22T20:13:58Z

+        return emb.mean(dim=1)
+
+
+class _PositionalEncoding(nn.Module):


it could go to layers/_encoders?

Actually its just a misinterpretation, it is an embedding not an encoder. We already have a PositionalEmbedding class in _embeddings/_positional_embedding.py doing the exact same math. Instead of duplicating that logic in an _encoder.py file, I created a _PositionalEmbedding child class inside the existing _positional_embedding.py file. It inherits the fixed sinusoidal buffer from the parent and adds the specific dropout and additive forward logic required for UniTS.

phoeenniixx · 2026-03-22T20:14:29Z

+        return self.drop(x + self.pe[:, : x.size(1), :])
+
+
+class _TransformerBlock(nn.Module):


should we add it to layer/_transforms or something?

I think creating a new _transforms folder wouldn't be necessary here because it would be confusing for new contributors as transform and transformers are 2 different terminologies. Furthermore, we already have an established layers/_blocks/ directory. A transformer block is standard neural network block logic, so keeping it grouped with other blocks in layers/_blocks/_transformer_block.py works perfectly entirely to our existing layout.

Muhammad-Rebaal · 2026-03-23T15:21:44Z

Hi @phoeenniixx ,
Kindly review the PR

Muhammad-Rebaal · 2026-06-30T18:10:59Z

Hi @phoeenniixx,

Kindly review this and let me know if we're done with this one so I'd add the api-reference for this one.

phoeenniixx · 2026-07-01T11:30:48Z

what is the use of this folder?

The model file already imports directly from the canonical locations. I'll remove the _units layer folder entirely.

phoeenniixx · 2026-07-01T11:31:57Z

+    )
+
+
+def test_basic_model_initialization(model, basic_metadata):


i think test framework is enough for these tests?

Yes, the extra tests are redundant I've removed that.

phoeenniixx

I think it is almost ready. Just please de-duplicate the tests that can already being tested using the test framework

…Class

phoeenniixx · 2026-07-21T08:13:07Z

+        "info:name": "UniTS",
+        "info:compute": 4,
+        "authors": ["Muhammad-Rebaal", "gasvn", "sohamukute"],
+        "python_dependencies": ["torch"],


why do we need this? isnt torch already a core dep here?

Actually I implemented this seeing the model extension v2. So I think I its an essential tag to be added torch as dependencies, not knowing the core-dep can't be added. Although I removed that.

phoeenniixx · 2026-07-21T08:13:55Z

+            {},
+            {
+                "patch_len": 8,
+                "stride": 4,


please add some loss functions as well here - Is this model only compatible with point prediction losses, or can it also handle quantile and distribution losses?

In the paper the model it is mentioned only about point prediction losses but we can extend it so I extend it to both quantile and distribution losses.

phoeenniixx · 2026-07-21T08:17:55Z

why do we have this file? I dont understand, it has just have a few fixtures and no clear tests are written,
Can you explain what this file is exactly doing?

Actually the model was initially implemented by the use of TslibBaseModel which later converted to EncoderDecoder so I removed those tests manually and new tests yet to be added which I added in my latest commit.

phoeenniixx

Please remove the test file, or add some tests there - currently it feels like an AI hallucination

Muhammad-Rebaal · 2026-07-27T14:52:26Z

Changes made in the recent commit:

removed the extra tag.
Added support of losses (quantile, distribution)
Write tests for the model.

phoeenniixx · 2026-07-28T14:42:01Z

+        return EncoderDecoderTimeSeriesDataModule
+
+    @classmethod
+    def get_base_test_params(cls):


As I mentioned earlier, get_base_test_params is not going to work for now for v2
use get_test_train_params instead.
for that you need to remove pred_type, y_type tags

looping over all the losses is not a good idea rn for v2

Aren't these tags pred_type, y_type be removed from here : https://github.com/sktime/pytorch-forecasting/blob/main/extension_templates/v2/model_simple/model_pkg.py if these are part of get_base_test_params ?

looping over all the losses is not a good idea rn for v2

Can you suggest me the preferred ones ? so I'd add those

Aren't these tags pred_type, y_type be removed from here : https://github.com/sktime/pytorch-forecasting/blob/main/extension_templates/v2/model_simple/model_pkg.py if these are part of get_base_test_params ?

I have reservations about this :)
Yes, we can remove this, but we will have to add it again, once the work on updating the test framework is done - which is the next work item on my list. As, we dont have much contributors around at this time, we can manage to ignore this. wdyt?

Muhammad-Rebaal · 2026-07-31T14:12:43Z

Hi @phoeenniixx ,
I've made the following changes as mentioned:

Remove the tags and renamed get_base_test_params to get_test_train_params

Also I add another change:

Removed the Distribution loss because I didn't find any implementation of the Distribution loss in base_model_v2 ? If there's present kindly let me know I'd implement it...

…ecasting into units_v2

phoeenniixx · 2026-08-08T07:22:18Z

+            instance. ``create_test_instance`` uses the first (or only) dictionary in
+            ``params``.
+        """
+        from pytorch_forecasting.metrics import NormalDistributionLoss, QuantileLoss


I am not sure if DistributionLoss works well with v2 rn, does it?

No I didn't find any implementation as well that's why I asked you above. That's why in my last commit I remove the code implementation as well but forget to remove this dead import sorry for confusion

Is there any open issue regarding the support of Distribution loss in v2 models ?

[ENH] Units_v2 Model added

930f855

Muhammad-Rebaal requested review from PranavBhatP, benHeid, fkiraly, fnhirwa, jdb78, phoeenniixx and yarnabrina as code owners March 9, 2026 19:37

Merge branch 'main' into units_v2

0f4faf4

sohamukute mentioned this pull request Mar 10, 2026

[ENH] add UniTS v2 model (#2158) #2163

Closed

3 tasks

phoeenniixx reviewed Mar 15, 2026

View reviewed changes

Comment thread pytorch_forecasting/models/units/_units_pkg_v2.py

phoeenniixx reviewed Mar 15, 2026

View reviewed changes

phoeenniixx added enhancement New feature or request module:models labels Mar 15, 2026

Muhammad-Rebaal and others added 2 commits March 16, 2026 17:24

Merge branch 'main' into units_v2

43dd89d

Add UniTS model, its package container, and integrate it into the mod…

2d46f32

…els initialization

Muhammad-Rebaal and others added 4 commits March 17, 2026 10:48

[ENH] Added a default fixture

e9cd22f

feat: Add UniTS_pkg_v2 for UniTS model definition, metadata, and te…

c8c81a9

…st parameters.

[BUG] Fix default fixture context_length and shared dict

61e4d29

Merge branch 'main' into units_v2

ee65d40

phoeenniixx reviewed Mar 22, 2026

View reviewed changes

Muhammad-Rebaal and others added 2 commits March 23, 2026 20:01

Merge branch 'main' into units_v2

7e13f3b

fix: Code Refactored

8900902

Muhammad-Rebaal and others added 2 commits June 29, 2026 14:32

Merge branch 'main' into units_v2

4327f59

fix: removed the v1 tags

1fa26b3

phoeenniixx reviewed Jul 1, 2026

View reviewed changes

phoeenniixx requested changes Jul 1, 2026

View reviewed changes

Muhammad-Rebaal and others added 4 commits July 3, 2026 16:38

fix: Performed De-duplication

8e7919f

Added API reference for the Units model

7eb7fce

feat: Converted the model on the BaseClass instead of the TslibModel …

0dc7a8d

…Class

Merge branch 'main' into units_v2

26e8290

phoeenniixx reviewed Jul 21, 2026

View reviewed changes

phoeenniixx requested changes Jul 21, 2026

View reviewed changes

Muhammad-Rebaal and others added 2 commits July 26, 2026 13:34

Merge branch 'main' into units_v2

013bd15

remove the non-important tag, losses suport added, and tests added

96c229e

Muhammad-Rebaal added 2 commits July 27, 2026 22:51

fix pytest error

36f2e6a

fix code quality

dccfe8c

phoeenniixx reviewed Jul 28, 2026

View reviewed changes

Muhammad-Rebaal added 2 commits July 30, 2026 21:05

fix: tag removed and renamed param name

3c3b662

remove dist loss as there isn't any implementation in base_model_v2

e069a29

Muhammad-Rebaal added 3 commits August 3, 2026 17:00

Merge branch 'main' of https://github.com/Muhammad-Rebaal/pytorch-for…

874e477

…ecasting into units_v2

fix : Code Quality

0de97fc

fix Code Quality

038d8c4

phoeenniixx reviewed Aug 8, 2026

View reviewed changes

Muhammad-Rebaal and others added 3 commits August 8, 2026 14:14

fix : Remove dead import

30b0aab

Merge branch 'main' into units_v2

aca98b1

Merge branch 'main' into units_v2

81fc45a

		return self.drop(x + self.pe[:, : x.size(1), :])


		class _TransformerBlock(nn.Module):

		)


		def test_basic_model_initialization(model, basic_metadata):

Uh oh!

Conversation

Muhammad-Rebaal commented Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Mar 16, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoeenniixx Mar 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Mar 23, 2026

Uh oh!

Muhammad-Rebaal commented Jun 30, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Jul 27, 2026

Changes made in the recent commit:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Jul 31, 2026

Muhammad-Rebaal commented Mar 9, 2026 •

edited

Loading

codecov Bot commented Mar 9, 2026 •

edited

Loading

phoeenniixx Mar 22, 2026 •

edited

Loading

Muhammad-Rebaal Jul 28, 2026 •

edited

Loading