refactor(autojac): Avoid concat in jac_to_grad#515
Closed
PierreQuinton wants to merge 5 commits intomainfrom
Closed
refactor(autojac): Avoid concat in jac_to_grad#515PierreQuinton wants to merge 5 commits intomainfrom
PierreQuinton wants to merge 5 commits intomainfrom
Conversation
- moves Matrix and PSDMatrix to compute_gramian (not best position probably, but should be in _utils) - Change return type of compute_gramian to PSDMatrix - Add compute_gramian_sum (note that the responsability of casting to PSDMatrix is given to _utils now). - add _gramian_based version of jac_to_grad. Note that we could put the tensordot(weights, jacobian, dims=1) in _utils as a weight_generalize_matrix method.
PierreQuinton
added
cc: feat
Codecov Report✅ All modified and coverable lines are covered by tests.
🚀 New features to boost your workflow:
|
PierreQuinton
commented
Jan 15, 2026
| grad_outputs: tuple[Tensor, ...], | ||
| args: tuple[PyTree, ...], | ||
| kwargs: dict[str, PyTree], | ||
| ) -> Optional[Tensor]: |
Contributor
Author
There was a problem hiding this comment.
This could be a Optional[PSDMatrix]
ValerianRey
changed the title
ValerianRey
deleted the
avoid-concatenation-for-gramian-based-aggregators
branch
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The goal is to half peak memory usage when the aggregator is Gramian based.
Does:
_utilspackage totorchjdMatrixandPSDMatrixin_utilspackage_utils.compute_gramian: Tensor -> PSDMatrix_utils.compute_gramian_sum(Note that _utils is the only package responsible for cast toPSDMatrix)_jacobian_based(current) and_gramian_based(new) strategy tojac_to_grad@ValerianRey This could be cut into several more atomic PRs, this is why this is a draft (I just wanted to see what are the things/problems involved in this).