Add a Conv2D layer, and stop its gradient going through a batched matmul - #99
Merged
Merged
Conversation
Measured 1.5-3.2x on the kernel gradient across CNN-shaped layers; the forward is unchanged, since the rewriter already normalizes both forms to the same graph.
Codecov Report✅ All modified and coverable lines are covered by tests.
Additional details and impacted files@@ Coverage Diff @@
## main #99 +/- ##
==========================================
- Coverage 97.38% 89.95% -7.44%
==========================================
Files 56 56
Lines 2604 2609 +5
==========================================
- Hits 2536 2347 -189
- Misses 68 262 +194 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Nothing downstream could size itself from a conv stack before this: every output extent came back None even when the input was fully known.
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.
We had
Conv1Dand no 2-D convolution.Conv2Disn_spatial = 2and a docstring, because the layer base and the op under it were already rank-generic — so the substance here is what rank 2 turned up.The kernel gradient was going through a batched contraction instead of one GEMM.
_correlatecontracted a 4-D patch tensor against a 2-D kernel, andpt.gradreads the graph as written, before rewrites normalize the forward. Collapsing the batch and window axes first makes the gradient two plain matmuls — 1.5–3.2x ondW, forward unchanged. Separately, every backend dispatch test ran at rank 1, so a 2-D convolution reached jax, mlx and torch with nothing checking the result. Each now has a rectangular-kernel case, which is what catches a transposed spatial axis.