vect rrule pullbacks index a thunked cotangent once per argument (quadratic), same pattern as #860 #862

Follow-up to #860 / #861. The rrules for Base.vect (src/rulesets/Base/array.jl, three methods) read the input cotangent once per argument (projects[n](ȳ[n]), ȳ[n], NTuple{N}(ȳ)). When ȳ is a thunk, each access re-evaluates the thunk, so the pullback is quadratic in the number of arguments. This is the same pattern as the reduce(vcat/hcat) pullbacks fixed in #861.

using ChainRules, ChainRulesCore
for N in (256, 1024, 4096)
  X = ntuple(i -> iseven(i) ? 1.0 : 1.0f0, N)      # hits the Union{Number,AbstractArray{<:Number}} method
  y, back = rrule(Base.vect, X...)
  dY = @thunk(ones(length(y)))
  back(dY); back(unthunk(dY))
  t1 = minimum(@elapsed(back(dY)) for _ in 1:20); t2 = minimum(@elapsed(back(unthunk(dY))) for _ in 1:20)
  println("N=$N  thunk: ", round(t1*1e6, digits=1), " µs   unthunked: ", round(t2*1e6, digits=1), " µs")
end

At N = 2048 all three methods (mixed numbers, Any, and uniform Float64) were measured at roughly 600 to 670 µs with a thunk against roughly 40 µs after unthunk. ChainRules 1.74.8, ChainRulesCore 1.26.8, Julia 1.13.2.

Proposal: ȳ = unthunk(ȳ) at the top of each vect_pullback. Downstream in Zygote the cotangent arriving from sum is a thunk, as in FluxML/Zygote.jl#1686.

🧔‍♂️🤖 Assisted by AI