Skip to content

feat(aggregation) Add scale_mode parameter to AlignedMTL#527

Merged
ValerianRey merged 5 commits intoSimplexLab:mainfrom
rkhosrowshahi:aligned-mtl-scale-mode
Jan 23, 2026
Merged

feat(aggregation) Add scale_mode parameter to AlignedMTL#527
ValerianRey merged 5 commits intoSimplexLab:mainfrom
rkhosrowshahi:aligned-mtl-scale-mode

Conversation

@rkhosrowshahi
Copy link
Copy Markdown
Contributor

@rkhosrowshahi rkhosrowshahi commented Jan 23, 2026

Summary

  • Add a scale_mode parameter to AlignedMTL and AlignedMTLWeighting.
  • Support min (default), median, and rmse scaling for the balance transformation.

Tests

  • UV_NO_SYNC=1 uv run pytest tests/unit (2604 passed, 63 skipped, 33 xfailed)

Notes

This PR supersedes #526, which was auto-closed after the head repository was deleted (by mistake, sorry).

rkhosrowshahi and others added 4 commits January 23, 2026 10:12
Expose median and rmse scaling modes for the balance transformation to match
original behavior while keeping min as default.
Add coverage for median/rmse modes and invalid scale_mode handling.
Document the new scale_mode parameter for AlignedMTL and AlignedMTLWeighting.
@rkhosrowshahi
Copy link
Copy Markdown
Contributor Author

CI checks may show as pending until a maintainer approves workflows for this first-time contributor fork. Local unit tests passed (see PR body).

@codecov
Copy link
Copy Markdown

codecov bot commented Jan 23, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines Coverage Δ
src/torchjd/aggregation/_aligned_mtl.py 100.00% <100.00%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ValerianRey
Copy link
Copy Markdown
Contributor

Thanks for the PR, it's really clean and ready to merge now!

I just changed tests slightly, otherwise some of them fail when testing on cuda (on my laptop's GeForce GTX 1080 gpu, something that the CI cannot test automatically).

@ValerianRey ValerianRey changed the title Add scale_mode options to AlignedMTL feat(aggregation) Add scale_mode parameter to AlignedMTL Jan 23, 2026
@ValerianRey ValerianRey added cc: feat Conventional commit type for new features. package: aggregation labels Jan 23, 2026
@ValerianRey ValerianRey merged commit 8653929 into SimplexLab:main Jan 23, 2026
14 of 15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cc: feat Conventional commit type for new features. package: aggregation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants