Skip to content

Unthunk the cotangent in reduce(vcat/hcat) and hvcat pullbacks - #861

Open
CarloLucibello wants to merge 1 commit into
JuliaDiff:mainfrom
CarloLucibello:cl/unthunk-reduce-cat
Open

CarloLucibello wants to merge 1 commit into
JuliaDiff:mainfrom
CarloLucibello:cl/unthunk-reduce-cat

Conversation

@CarloLucibello

Copy link
Copy Markdown
Member

Closes #860.

The pullbacks of reduce(vcat, As), reduce(hcat, As) (reduce_vcat_pullback, reduce_hcat_pullback_1, reduce_hcat_pullback_2) and hvcat (hvcat_pullback) index the cotangent once per input. Indexing an AbstractThunk calls unthunk each time, so a thunked cotangent (e.g. from sum's pullback in Zygote, FluxML/Zygote.jl#1686) is recomputed length(As) times. This calls unthunk once at the top of each, as hcat/vcat/cat/stack already do.

I used dY = unthunk(ȳ) (renaming the argument) rather than dY = unthunk(dY), matching the neighbouring pullbacks and avoiding rebinding a variable captured by the do blocks. hvcat_pullback has the same pattern and is included.

Measured on Julia 1.13.2 (min of 5 runs), thunked cotangent @thunk(ones(size(Y))):

pullback before after plain array
reduce(vcat), 1600 × Vector(10) 3.9 ms, 196 MiB 0.026 ms, 0.35 MiB 0.024 ms
same, InplaceableThunk 5.8 ms, 196 MiB 0.030 ms, 0.35 MiB 0.023 ms
reduce(vcat), 800 × Matrix(10×2) 1.5 ms, 98 MiB 0.026 ms, 0.31 MiB 0.017 ms
reduce(hcat), 1600 × Vector(10) 4.3 ms, 196 MiB 0.024 ms, 0.35 MiB 0.024 ms
reduce(hcat), 800 × Matrix(10×2) 1.6 ms, 98 MiB 0.025 ms, 0.31 MiB 0.020 ms
hvcat, 200 × Vector(10) 0.145 ms, 3.1 MiB 0.010 ms, 0.05 MiB 0.003 ms

Added a testset in test/rulesets/Base/array.jl checking that @thunk and InplaceableThunk cotangents 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 about length(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/Base on 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:

  • vect pullbacks 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.
  • Other pullbacks I looked at in src/rulesets/Base/ either already unthunk or access the cotangent once.

🧔‍♂️🤖 Assisted by AI

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

codecov Bot commented Oct 31, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.64%. Comparing base (a441bd5) to head (d32589b).

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

reduce(vcat)/reduce(hcat) pullbacks index a thunked cotangent per slice without unthunk (quadratic)

1 participant