[ENH] Adding Support of Distribution Loss in v2 Models by Muhammad-Rebaal · Pull Request #2420 · sktime/pytorch-forecasting

Muhammad-Rebaal · 2026-09-09T13:32:47Z

Hi @phoeenniixx,
In order to add support I've done the following changes :

Added transform_output in BaseModel to scale predictions back to their real-world numbers.
Added output_size in BaseModel to tell the model how many numbers to predict.
Added target_normalizer to __init__ in TslibBaseModel to pass the data scaler to the model.

codecov · 2026-09-09T21:58:08Z

Codecov Report

❌ Patch coverage is 88.05970% with 8 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@0acb5a0). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
...ch_forecasting/models/base/_tslib_base_model_v2.py	25.00%	3 Missing ⚠️
...h_forecasting/models/tide/_tide_dsipts/_tide_v2.py	80.00%	3 Missing ⚠️
pytorch_forecasting/models/base/_base_model_v2.py	94.28%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2420   +/-   ##
=======================================
  Coverage        ?   88.90%           
=======================================
  Files           ?      218           
  Lines           ?    11713           
  Branches        ?        0           
=======================================
  Hits            ?    10414           
  Misses          ?     1299           
  Partials        ?        0

Flag	Coverage Δ
cpu	`88.90% <88.05%> (?)`
pytest	`88.90% <88.05%> (?)`

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

Can you please add some tests here? To see if the shapes are correct etc?

Maybe we can update the test framework even here - to loop over the losses (like in v1). That can happen in a separate PR stacked over this PR.

Muhammad-Rebaal · 2026-09-17T10:13:57Z

Hi @phoeenniixx,
I've added the tests

phoeenniixx · 2026-09-17T10:16:54Z

+def test_transform_output_dict_target_scale():
+    """Dict target_scale applies affine denormalization correctly."""
+    model = _make_model()
+    raw = torch.randn(4, 12, 1)
+    center = torch.tensor([10.0, 20.0, 30.0, 40.0])
+    scale = torch.tensor([2.0, 3.0, 4.0, 5.0])
+    target_scale = {"center": center, "scale": scale}
+
+    result = model.transform_output(raw, target_scale)
+
+    assert result.shape == raw.shape
+    assert torch.allclose(result[0], raw[0] * 2.0 + 10.0)
+    assert torch.allclose(result[3], raw[3] * 5.0 + 40.0)
+
+
+def test_transform_output_plain_tensor_target_scale():


how are these two tests useful? Why not direclty use the DistributionLoss?

These two tests only validate the affine denormalization path (pred * scale + center), which is the straightforward case.
I'll add those distribution loss tests as well

My doubt is - DistributionLoss does somethign similar no?
(i have not looked at the math of the losses, so i may be wrong)

These 2 tests are checking specifically making sure that Quantile Loss and Point Forecasts (like MAE) get un-scaled correctly, since I made changes in the transform_output function as well.

phoeenniixx

I think we are still lacking the tests that check If DistributionLoss works as intended with v2.

Also, this should be present in the test framework? why the base model?

By adding "support" we should also update the models which have this todo, no?

Muhammad-Rebaal · 2026-09-17T13:55:46Z

I think we are still lacking the tests that check If DistributionLoss works as intended with v2.

Ok I'll add that specific tests as well

Also, this should be present in the test framework? why the base model?

Yeah I got your point up here so that each time it would run it would also check it .

By adding "support" we should also update the models which have this todo, no?

I thought for the model we'd open an umbrella issue that would adding support for each. Although If that's your vision to add in all models I'll update that as well.

phoeenniixx · 2026-09-17T13:57:04Z

I thought for the model we'd open an umbrella issue that would adding support for each. Although If that's your vision to add in all models I'll update that as well.

We should try with atleast 1-2 models otherwise how would we know if this is working or not?

…l/pytorch-forecasting into distribution_loss

Muhammad-Rebaal · 2026-09-25T15:22:01Z

Hi @phoeenniixx,

I've made these changes :

Added support for 2 models TiDe and UniTS
Test added in test test_base_model_v2.py because we made changes in base_model_v2.py, test for specific models will auto-detected by the test_all_v2 framework via the NormalDistributionLoss() configs in the pkg files?

phoeenniixx · 2026-10-04T19:12:53Z

                "max_prediction_length": self.max_prediction_length,
                "min_encoder_length": self._min_encoder_length,
                "min_prediction_length": self._min_prediction_length,
+                "target_normalizer": self.target_normalizer,


should we pass a whole object in metadata? is this actually necessary? I think metadata should be lightweight and only should have the "metadata" and not objects unless it is actually 100% necessary

phoeenniixx · 2026-10-04T19:14:47Z

why are we changing the forward logic of this model?

There should be a cleaner and more robust way to handles the shapes?

phoeenniixx

Nice!
This is looking better.
I just have few concerns (pls see above)

Implemented Support for Distribution Loss

dac3f96

Muhammad-Rebaal requested review from benHeid, fkiraly, jdb78 and phoeenniixx as code owners September 9, 2026 13:32

Muhammad-Rebaal marked this pull request as draft September 9, 2026 14:35

Muhammad-Rebaal and others added 2 commits September 10, 2026 00:12

Fix: Namespace error

4cb6fdf

Merge branch 'main' into distribution_loss

641e6c1

Muhammad-Rebaal marked this pull request as ready for review September 10, 2026 13:34

phoeenniixx assigned Muhammad-Rebaal Sep 10, 2026

phoeenniixx requested changes Sep 10, 2026

View reviewed changes

phoeenniixx added enhancement New feature or request module:models module:metrics labels Sep 10, 2026

Muhammad-Rebaal and others added 2 commits September 17, 2026 00:37

Merge branch 'main' into distribution_loss

8caf327

Added tests

3983b4d

Merge branch 'main' into distribution_loss

fcd6d85

Muhammad-Rebaal requested a review from phoeenniixx September 17, 2026 10:15

phoeenniixx reviewed Sep 17, 2026

View reviewed changes

phoeenniixx requested changes Sep 17, 2026

View reviewed changes

Muhammad-Rebaal marked this pull request as draft September 23, 2026 21:40

Muhammad-Rebaal and others added 5 commits September 24, 2026 02:43

Added Distribution Loss support in Tide and UniTS models

95886f8

Merge branch 'distribution_loss' of https://github.com/Muhammad-Rebaa…

c0fbb2f

…l/pytorch-forecasting into distribution_loss

Merge branch 'main' into distribution_loss

333cb6c

Fix: Code Quality

2682958

Merge branch 'distribution_loss' of https://github.com/Muhammad-Rebaa…

7c1b947

…l/pytorch-forecasting into distribution_loss

Fix: Code Quality

fdb69dd

github-advanced-security Bot found potential problems Sep 24, 2026

View reviewed changes

Comment thread tests/test_models/test_base_model_v2.py Fixed

Muhammad-Rebaal added 2 commits September 24, 2026 19:29

Fix: Unknown import

0c8c615

Add test for Distribution loss

6bcaae4

Muhammad-Rebaal marked this pull request as ready for review September 25, 2026 15:06

phoeenniixx reviewed Oct 4, 2026

View reviewed changes

phoeenniixx requested changes Oct 4, 2026

View reviewed changes

Merge branch 'main' into distribution_loss

2aa4ad0

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal commented Sep 17, 2026

Uh oh!

phoeenniixx Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

phoeenniixx Sep 17, 2026

Choose a reason for hiding this comment

Uh oh!

Muhammad-Rebaal Sep 25, 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!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

phoeenniixx commented Sep 17, 2026

Uh oh!

Uh oh!

Muhammad-Rebaal commented Sep 25, 2026

Uh oh!

phoeenniixx Oct 4, 2026

Choose a reason for hiding this comment

Uh oh!

phoeenniixx Oct 4, 2026

Choose a reason for hiding this comment

Uh oh!

phoeenniixx Oct 4, 2026

Choose a reason for hiding this comment

Uh oh!

phoeenniixx left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Muhammad-Rebaal commented Sep 9, 2026 •

edited

Loading

codecov Bot commented Sep 9, 2026 •

edited

Loading

Muhammad-Rebaal Sep 25, 2026 •

edited

Loading

Muhammad-Rebaal commented Sep 17, 2026 •

edited

Loading