[Feature] Allow custom RINorm kwargs & fix foundation compatibility

[daidahao]

Checklist before merging this PR:

• Mentioned all issues that this PR fixes or addresses.
• Summarized the updates of this PR under Summary.
• Added an entry under Unreleased in the Changelog.

Fixes #3028, fixes #2748.

Summary

• All torch models now accept use_reversible_instance_norm as a flexible parameter that can be set to True, False, or a dictionary of RINorm hyperparameters. This allows users to customize RINorm for their models.
• Disallowed use_reversible_instance_norm=True or {"affine": True} for foundation models to prevent checkpoint loading errors due to incompatible weights.

Other Information

#2758 is not being fixed in this PR.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 95.66%. Comparing base (3ae6ca3) to head (0c41992).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #3029      +/-   ##
==========================================
- Coverage   95.73%   95.66%   -0.07%     
==========================================
  Files         158      158              
  Lines       17122    17131       +9     
==========================================
- Hits        16391    16389       -2     
- Misses        731      742      +11     

☔ View full report in Codecov by Sentry.
📢 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.

[dennisbader]

Thanks for the PR @daidahao 🚀

Could we rather under the hood allow PLForecastingModule to already accept a bool or any kwargs (e.g. {"affine": True}) to instantiate the RINorm? If it's a bool, it would use the existing behavior.

Then for the foundation models we would pass {"affine": False} whenever use_reversible_instance_norm=True.

This would be a quick fix while still allowing the foundation models to use RINorm (could be good for fine-tuning).

[daidahao]

@dennisbader

This would be a quick fix while still allowing the foundation models to use RINorm (could be good for fine-tuning).

Personally, I would not recommend using RINorm for prediction or fine-tuning since almost all foundation models have their own RINorm. That's why I have now removed any reference to RINorm in Chronos-2 and TimesFM. Nevertheless, in FoundationModel, use_reversible_instance_norm is being handled properly to produce warnings and override "affine" key to False.

For all torch models, I've updated PLForecastingModule so that it now accepts use_reversible_instance_norm as either a bool or dict of hyperparameters. One question remains, should we update all torch model doctorings to highlight this new feature?

Let me know what you think.

[dennisbader]

Looks great, thanks a lot @daidahao 🚀
I applied some minor updates (shortened a bit the parameter description and added it to all model docs)