feat(autojac): Make jac_to_grad return optional weights#586
Open
ValerianRey wants to merge 11 commits intomainfrom
Open
feat(autojac): Make jac_to_grad return optional weights#586ValerianRey wants to merge 11 commits intomainfrom
ValerianRey wants to merge 11 commits intomainfrom
Conversation
* Change `aggregator: Aggregator` to `method: Aggregator | Weighting` and return type to optional `Tensor`. * Make `method` positional only. * Add overloads to rename `method` to `aggregator` or `weighting` and link it to output type. * Compute the weights if we provide a weighting and return them. * Update the doc and add a usage example
ValerianRey
added
cc: feat
ValerianRey
force-pushed
the
make-jac-to-grad-return-weights-2
branch
from
5c79901 to
b701c59
Compare
- We never use it anywhere else in the library, and I think it's quite redundant here
PierreQuinton
approved these changes
Feb 21, 2026
Contributor
PierreQuinton
left a comment
PierreQuinton
left a comment
There was a problem hiding this comment.
Very nice, LGTM, but might require an update of changelog to specify this better.
ValerianRey
commented
Feb 21, 2026
Comment on lines
+108
to
+109
| weights = aggregator.weighting(jacobian_matrix) | ||
| gradient_vector = weights @ jacobian_matrix |
Contributor
There was a problem hiding this comment.
In this case, the hooks of the aggregator are not run. Not sure about how to resolve this without running the aggregator.
Contributor
There was a problem hiding this comment.
It feels like we could cheat a bit and have a boolean field return_weights to forward in WeightedAggregators that we don't expose, on True, it would then also return the weights, and run the hooks. There is probably a better solution.
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.
This is based on #585 but instead of taking a
methodargument that can be eitherAggregatororWeighting, it always takes anaggregator.jac_to_gradreturn optional weightsWeightingjac_to_gradonly take aggregatorThere's just one problem: currently,
WeightedAggregatoris not public. So we can't really talk about this class in the docstring ofjac_to_grad. I think we have two options:WeightedAggregatorpublic. Because it has aweightingparameter, of typeWeighting[Matrix], we'll have to find a solution for that too (makeMatrixpublic or annotateweightingasWeighting).Weighting" instead of "If a :class:WeightedAggregator <torchjd.aggregation._aggregator_bases.WeightedAggregator>is provided"A similar issue will arise when we'll add the option to do the gramian-sum optimization. We'll have to talk about
GramianWeightedAggregatorand make it public (as well as PSDMatrix), or we'll have to say "aggregator uses a gramian-based Weighting`.Note that this latter problem would also happen to some extent in #585, because
Weighting[PSDMatrix]is not public, so we'll also have to say something likeif weighting is gramian-basedor makePSDMatrixpublic.I think the easiest solution for now is just to talk about "aggregator using a
Weighting", and later about "aggregator using a gramian-basedWeighting" when we add the gramian-sum optimization. (Done in b701c59)