[ENH] Add TSMixer v2 model by seaic-mac-murchadha · Pull Request #2387 · sktime/pytorch-forecasting

seaic-mac-murchadha · 2026-08-20T13:27:28Z

Reference Issues/PRs

For issue #1992 which requires implementing "models to the v2 interface". This PR adds the TSMixer implementation for the v2 interface.

#1992

What does this implement/fix? Explain your changes.

Implements TSMixer model for the v2 interface, including:

temporal and channel mixing blocks
point forecasting
QuantileLoss support
model and input handling tests

What should a reviewer concentrate their feedback on?

The implementation is complete as an initial functional TSMixer v2 integration.

Quantile forecasting with multiple targets and inputs without target history will raise explicit errors. Feedback on methodology for resolving these cases without errors would be welcome if their support is required.

Did you add any tests for the change?

Yes, PR has corresponding tests implemented therein, covering:

model initialization
point forecasting
quantile forecasting
input preparation
error handling

Any other comments?

Only single-target quantile forecasting is currently supported, consistent with the current DLinear implementation.

PR checklist

The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [BUG] - bugfix, [MNT] - CI, test framework, [ENH] - adding or improving code, [DOC] - writing or improving documentation or docstrings.
Added/modified tests
Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
To run hooks independent of commit, execute pre-commit run --all-files

codecov · 2026-08-20T14:18:15Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@36ac67f). Learn more about missing BASE report.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2387   +/-   ##
=======================================
  Coverage        ?   88.22%           
=======================================
  Files           ?      199           
  Lines           ?    11142           
  Branches        ?        0           
=======================================
  Hits            ?     9830           
  Misses          ?     1312           
  Partials        ?        0

Flag	Coverage Δ
cpu	`88.22% <100.00%> (?)`
pytest	`88.22% <100.00%> (?)`

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.

…ided

phoeenniixx

Thanks!
Please add docstrings to the public classes and methods

Also, for the test file, i think most of those tests are already handled by the unified test framework no? Can you please de-duplicate the file

Also, can you please add the API reference for the model and pkg class as well?

phoeenniixx · 2026-08-22T18:50:00Z

+            class.
+        """
+
+        import torch.nn as nn


do we need this import?

I’m using it for nn.MSELoss() within the test. Would it be preferable to remove this case, or perhaps move the import to the top of the file?

Thanks, I missed it the last time

phoeenniixx · 2026-08-22T18:50:24Z

+from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel
+
+
+class TSMixerBlock(nn.Module):


Can you please add docstrings here?

Definitely, have added docstrings to the TSMixerBlock class and its forward function in the latest commit.

phoeenniixx · 2026-08-22T18:51:36Z

+            metadata=metadata,
+        )
+
+        warnings.warn(


I think the warning is already handled by BaseModel no? See here

Thanks @phoeenniixx, have removed the redundant warning now.

phoeenniixx · 2026-08-22T18:54:00Z

+    assert model.n_quantiles is None
+
+
+def test_tsmixer_forward(sample_dataset):


I think it is already being tested by the test framework

Yes, it isn't adding a meaningful test, have removed it now. Thanks @phoeenniixx.

seaic-mac-murchadha · 2026-08-24T20:06:07Z

Hi @phoeenniixx,

Thank you for the review. Have attempted to address the four comments left. Additionally, the generic tests have been removed and TSMixer has been added to the API documentation.

seaic-mac-murchadha · 2026-08-24T21:54:20Z

Hi @phoeenniixx,

The failure appears to be due to a network/DNS issue whilst downloading stallion.parquet, rather than the PR code changes. Would a re-run be possible?

phoeenniixx · 2026-08-27T10:58:45Z

+from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel
+
+
+class TSMixerBlock(nn.Module):


should this be added to layers/_blocks? Sorry i didnt notice this earlier

phoeenniixx · 2026-08-27T11:00:06Z

+
+
+@pytest.fixture
+def sample_dataset():


you can use examples.load_toydata as well here and then pass that to TimeSeries

phoeenniixx

Thanks! i think it is almost ready!
Just a few nit-picks

seaic-mac-murchadha added 17 commits August 20, 2026 10:26

[TSMixer v2] add init

33e5906

[TSMixer v2] add TSMixerBlock

beca811

[TSMixer v2] add unstable note

a8fd17b

[TSMixer v2] add TSMixer with init

20ecc8d

[TSMixer v2] add docstring for TSMixer

bda6af3

[TSMixer v2] for TSMixer class add _init_network()

6c91fcb

[TSMixer v2] for TSMixer class add _encoder()

cfcdf35

[TSMixer v2] for TSMixer class add _prepare_input_data()

f62c902

[TSMixer v2] for TSMixer class add forward()

0f35a8d

[TSMixer v2] add package configuration

617f1e5

[TSMixer v2] support QuantileLoss

ac7990a

[TSMixer v2] for testing add sample dataset

3a17d25

[TSMixer v2] add model tests

e2ce3ec

[TSMixer v2] update _prepare_input_data()

b0338c7

[TSMixer v2] add testing for _prepare_input_data()

d48bb62

[TSMixer v2] add testing for error handling

8e0c437

[TSMixer v2] apply ruff formatting

ba79905

seaic-mac-murchadha requested review from benHeid, fkiraly, jdb78 and phoeenniixx as code owners August 20, 2026 13:27

[TSMixer v2] add test for forward prediction with target scaling prov…

f9d4837

…ided

phoeenniixx requested changes Aug 22, 2026

View reviewed changes

phoeenniixx added enhancement New feature or request module:models ptf-v2 Related to `pytorch-forecasting` v2 labels Aug 22, 2026

seaic-mac-murchadha force-pushed the tsmixer-v2 branch 2 times, most recently from 8277cc0 to 9b7ab27 Compare August 23, 2026 22:03

Add docstrings to TSMixerBlock

d2e9470

seaic-mac-murchadha force-pushed the tsmixer-v2 branch from 9b7ab27 to d2e9470 Compare August 23, 2026 22:05

seaic-mac-murchadha added 4 commits August 23, 2026 23:06

Remove redundant experimental model warnings

af66d3f

Remove redundant TSMixer forward test

d79be74

Add TSMixer v2 API reference for the model and pkg

3b8e43c

Remove non-model specific tests for TSMixer v2

ca2f59d

Merge branch 'main' into tsmixer-v2

1ea1969

Merge branch 'main' into tsmixer-v2

336d164

phoeenniixx reviewed Aug 27, 2026

View reviewed changes

phoeenniixx requested changes Aug 27, 2026

View reviewed changes

		from pytorch_forecasting.models.base._tslib_base_model_v2 import TslibBaseModel


		class TSMixerBlock(nn.Module):

		assert model.n_quantiles is None


		def test_tsmixer_forward(sample_dataset):

Uh oh!

Conversation

seaic-mac-murchadha commented Aug 20, 2026

Reference Issues/PRs

What does this implement/fix? Explain your changes.

What should a reviewer concentrate their feedback on?

Did you add any tests for the change?

Any other comments?

PR checklist

Uh oh!

codecov Bot commented Aug 20, 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 • 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!

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!

seaic-mac-murchadha commented Aug 24, 2026

Uh oh!

seaic-mac-murchadha commented Aug 24, 2026

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Aug 20, 2026 •

edited

Loading

phoeenniixx left a comment •

edited

Loading