Repository navigation
Unthunk the cotangent in reduce(vcat/hcat) and hvcat pullbacks - #861
Open
CarloLucibello wants to merge 1 commit into
Open
CarloLucibello wants to merge 1 commit into
CarloLucibello wants to merge 1 commit into
Conversation
These pullbacks index the cotangent once per input array. Indexing an `AbstractThunk` forwards to `unthunk`, which recomputes the whole array on every access, so the cost was O(length(As) * length(dY)). Unthunk once at the top of each pullback, as the other `cat` pullbacks already do. Also bumps the patch version. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #861 +/- ##
=======================================
Coverage 96.63% 96.64%
=======================================
Files 28 28
Lines 3598 3602 +4
=======================================
+ Hits 3477 3481 +4
Misses 121 121 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This branch has not been deployed
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
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.
Closes #860.
The pullbacks of
reduce(vcat, As),reduce(hcat, As)(reduce_vcat_pullback,reduce_hcat_pullback_1,reduce_hcat_pullback_2) andhvcat(hvcat_pullback) index the cotangent once per input. Indexing anAbstractThunkcallsunthunkeach time, so a thunked cotangent (e.g. fromsum's pullback in Zygote, FluxML/Zygote.jl#1686) is recomputedlength(As)times. This callsunthunkonce at the top of each, ashcat/vcat/cat/stackalready do.I used
dY = unthunk(ȳ)(renaming the argument) rather thandY = unthunk(dY), matching the neighbouring pullbacks and avoiding rebinding a variable captured by thedoblocks.hvcat_pullbackhas the same pattern and is included.Measured on Julia 1.13.2 (min of 5 runs), thunked cotangent
@thunk(ones(size(Y))):reduce(vcat), 1600 ×Vector(10)InplaceableThunkreduce(vcat), 800 ×Matrix(10×2)reduce(hcat), 1600 ×Vector(10)reduce(hcat), 800 ×Matrix(10×2)hvcat, 200 ×Vector(10)Added a testset in
test/rulesets/Base/array.jlchecking that@thunkandInplaceableThunkcotangents give the same result as the plain array for all the variants above (vectors / matrices / mixed), and that@allocated back(thunk) < 10 * sizeof(Y)at 1600 inputs (the quadratic path allocates aboutlength(As) * sizeof(Y)). Without the fix, the 10 allocation checks fail and the 14 equality checks pass; with it all 24 pass, also under--check-bounds=yes --code-coverage=user.Tests:
test/runtests.jl rulesets/Baseon Julia 1.13.2 passes (16978 pass, 52 broken as before). Not run locally on 1.10 (left to CI).Bumps the patch version to 1.74.10, as recent PRs do.
Noticed but left alone:
vectpullbacks indexȳ[n]per element and are equally quadratic in the number of arguments (N=256: 0.8 ms vs 3 µs), but N is a compile-time argument count, so I did not touch them.src/rulesets/Base/either alreadyunthunkor access the cotangent once.🧔♂️🤖 Assisted by AI