Skip to content

refactor(aggregation)!: Remove generalized gramians#692

Merged
ValerianRey merged 12 commits into
mainfrom
remove-generalized-gramians
May 22, 2026
Merged

refactor(aggregation)!: Remove generalized gramians#692
ValerianRey merged 12 commits into
mainfrom
remove-generalized-gramians

Conversation

@ValerianRey
Copy link
Copy Markdown
Contributor

@ValerianRey ValerianRey commented May 20, 2026

Closes #690

  • Engine.compute_gramian now always returns a flat [m, m] gramian, regardless of the output shape
  • Removed GeneralizedWeighting, and Flattening — they are no longer needed
  • Updated the IWMTL example to use UPGradWeighting directly and reshape the weights before calling backward
  • Added a migration guide in CHANGELOG.md

Simplify `Engine.compute_gramian` to always return a flat `[m, m]` PSD matrix
where `m = output.numel()`, removing the concepts of generalized Gramians,
`PSDTensor`, `GeneralizedWeighting`, and `Flattening`.

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 breaking-change This PR introduces a breaking change. package: autogram labels May 20, 2026
@github-actions github-actions Bot changed the title feat(autogram): Remove generalized gramians refactor!: Remove generalized gramians May 20, 2026
@ValerianRey ValerianRey force-pushed the remove-generalized-gramians branch from ce15075 to 5ad48d2 Compare May 20, 2026 23:59
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.

I think I like it but let me think about it. Also I don't think we want to erase all private generalized_gramian utilities and put their implementation in autogram.

Comment thread src/torchjd/_linalg/_gramian.py Outdated
Comment thread src/torchjd/_linalg/_generalized_gramian.py
Comment thread src/torchjd/autogram/_engine.py Outdated
@ValerianRey ValerianRey marked this pull request as ready for review May 21, 2026 20:26
@ValerianRey ValerianRey requested a review from a team as a code owner May 21, 2026 20:26
@ValerianRey
Copy link
Copy Markdown
Contributor Author

@PierreQuinton this is now ready for review. I reverted to (almost) the original implementation of autogram: compute square gramian, reshape, movedim. The extra flatten is because we now want to flatten.
I also reviewed everything.

In a future PR we should be able to let compute_gramian take any number of output tensors (instead of a single one), to match the interface of autojac. There are a bunch of other changes to make the interfaces of autogram and autojac closer together, but let's go step by step.

@github-actions github-actions Bot changed the title refactor!: Remove generalized gramians refactor(aggregation)!: Remove generalized gramians May 22, 2026
@PierreQuinton
Copy link
Copy Markdown
Contributor

In a future PR we should be able to let compute_gramian take any number of output tensors (instead of a single one), to match the interface of autojac. There are a bunch of other changes to make the interfaces of autogram and autojac closer together, but let's go step by step.

That would be nice. But I think that this is not trivial with the batched dims (I guess all could have one, and possibly a different? I think in principle that would work but not trivial).

Comment thread src/torchjd/autogram/_engine.py Outdated
Co-authored-by: Pierre Quinton <pierre.quinton@epfl.ch>
@ValerianRey ValerianRey merged commit 0698f40 into main May 22, 2026
15 checks passed
@ValerianRey ValerianRey deleted the remove-generalized-gramians branch May 22, 2026 12:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking-change This PR introduces a breaking change. 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.

Remove generalized gramians

2 participants