Skip to content

refactor(aggregation): Make NashMTL and CAGrad always importable#679

Merged
ValerianRey merged 3 commits into
mainfrom
refactor/optional-deps-mixin
May 12, 2026
Merged

refactor(aggregation): Make NashMTL and CAGrad always importable#679
ValerianRey merged 3 commits into
mainfrom
refactor/optional-deps-mixin

Conversation

@ValerianRey
Copy link
Copy Markdown
Contributor

@ValerianRey ValerianRey commented May 11, 2026

Summary

  • Add _WithOptionalDeps mixin to _mixins.py: subclasses declare _REQUIRED_DEPS and _INSTALL_HINT; __init__ raises ImportError with a clear install message if any dep is missing
  • CAGradWeighting and _NashMTLWeighting now inherit _WithOptionalDeps first, so the check fires at instantiation rather than at import time
  • CAGrad, CAGradWeighting, and NashMTL are now unconditionally importable and listed in __all__
  • Tests updated to use pytest.importorskip (the standard pattern) instead of a manual try/except ImportError guard
  • Add changelog entry

🤖 Generated with Claude Code

Add _WithOptionalDeps mixin that raises ImportError at instantiation time
when optional dependencies are missing, replacing the module-level guard
that previously prevented import altogether.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@ValerianRey ValerianRey requested review from a team and PierreQuinton as code owners May 11, 2026 15:34
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@ValerianRey ValerianRey added package: aggregation cc: refactor Conventional commit type for any refactoring, not user-facing, and not typing or perf improvements labels May 11, 2026
@ValerianRey
Copy link
Copy Markdown
Contributor Author

@PierreQuinton IMO this is ready. You can merge if it helps for the rest.

Copy link
Copy Markdown
Contributor

@PierreQuinton PierreQuinton left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. In the context of QP solvers and of #678 do you think we should make the aggregator or the projector a _WithOptionalDeps? I think that this would change where the mixin would go to.

Comment thread src/torchjd/aggregation/_mixins.py Outdated
Co-authored-by: Pierre Quinton <pierre.quinton@epfl.ch>
@ValerianRey
Copy link
Copy Markdown
Contributor Author

LGTM. In the context of QP solvers and of #678 do you think we should make the aggregator or the projector a _WithOptionalDeps?

Yes

I think that this would change where the mixin would go to.

Maybe we should just add a torchjd.aggregation.projector package instead of having the projectors in linalg. The reason we don't have Matrix and PSDMatrix in aggregation is that they're not only used in aggregation but also in autogram.

@ValerianRey ValerianRey merged commit 6edc33e into main May 12, 2026
15 checks passed
@ValerianRey ValerianRey deleted the refactor/optional-deps-mixin branch May 12, 2026 12:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cc: refactor Conventional commit type for any refactoring, not user-facing, and not typing or perf improvements package: aggregation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants