Closed
Conversation
Add reshape of jacobian for the scalar output case. Fix reshape of the Gramian, we for the last half of the dimensions, we need to reshape in the same order as the first, then we move the dimensions. We could in principle create a `reshape_gramian` function that does this, as well as a `move_dim_gramian` Add a test of values for all four cases of having a batched/non-batched dimension. Tests or reshape/move-dim should work should go in another test. Remove some tests that do not test anything more than `test_gramian_is_correct`. Add `_gramian_utils.py` which contains helper to `reshape` and `movedim` on a Gramian. Add `generate_vmap_rule = True` for `JacobianAccumulator`. This allows vmaping the forward phase. This enables having several Engines defined on the same module. Add `test_reshape_equivariance` Add tests to verify that gramian utils yields the correct quadratic forms. Add tests to verify that gramian utils yields the correct quadratic forms. Add `test_movedim_equivariance` Fix warning. Fix warning. Remove handles from `ModuleHookManager` Change `batched_dims` to a single optional `batched_dim`. Fix movedim in `compute_gramian` and add `test_movedim_equivariance` Remove `grad_output`, can be added later, but should be `jac_output` instead. Make modules with incompatible batched operations are compatible with non-batched autogram. Fix doc tests Provide the autograd vjp for when no dimension is batched. This enables having a single forward in that case which should be faster. Make VJPs into Callable classes.
Co-authored-by: Valérian Rey <31951177+ValerianRey@users.noreply.github.com>
Co-authored-by: Valérian Rey <31951177+ValerianRey@users.noreply.github.com>
Not the final name I think, but at least it's consistent with the method name
Co-authored-by: Pierre Quinton <pierre.quinton@epfl.ch>
At this point, 3 architectures fail: SomeFrozenParam, SomeUnusedParam and MultiOutputWithFrozenBranch
… variables This fixes non-batched engien on SomeFrozenParams architecture
…grad in AutogradVJP This fixes non-batched engine on SomeUnusedParam
…adVJP This fixes non-batched engine on MultiOutputWithFrozenBranch
Maybe not a definitive name, but I think it's more clear
* Small improvement of clarity
…e can also contain (at most) one element set to -1, the size of that dimension is deduced from the total number of elements
Contributor
Author
|
Closing this PR in favor of an archive branch (archive/add-hierarchical-weighting). We can start a new PR when we want to experiment on this. |
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 PR contains 0a25555 which adds the
HierarchicalWeighting.I think we should wait to test this on proper experiments / speed benchmarks, before merging.