diff --git a/Project.toml b/Project.toml index a18c5c4b0..952b846d4 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.9" +version = "1.74.10" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Base/array.jl b/src/rulesets/Base/array.jl index 7e4098509..9b12d1439 100644 --- a/src/rulesets/Base/array.jl +++ b/src/rulesets/Base/array.jl @@ -50,7 +50,8 @@ function rrule( ) where {N} projects = map(ProjectTo, X) function vect_pullback(ȳ) - X̄ = ntuple(n -> projects[n](ȳ[n]), N) + dY = unthunk(ȳ) + X̄ = ntuple(n -> projects[n](dY[n]), N) return (NoTangent(), X̄...) end return Base.vect(X...), vect_pullback @@ -58,7 +59,10 @@ end # Data is unmodified, so no need to project. function rrule(::typeof(Base.vect), X::Vararg{Any,N}) where {N} - vect_pullback(ȳ) = (NoTangent(), ntuple(n -> ȳ[n], N)...) + function vect_pullback(ȳ) + dY = unthunk(ȳ) + return (NoTangent(), ntuple(n -> dY[n], N)...) + end return Base.vect(X...), vect_pullback end @@ -259,12 +263,14 @@ function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVe Y = reduce(hcat, As) Base.require_one_based_indexing(Y) widths = map(A -> size(A,2), As) - function reduce_hcat_pullback_2(dY) + project_As = ProjectTo.(As) + function reduce_hcat_pullback_2(ȳ) + dY = unthunk(ȳ) hi = Ref(0) - dAs = map(widths) do w + dAs = map(project_As, widths) do project, w lo = hi[]+1 hi[] += w - dY[:, lo:hi[]] + project(dY[:, lo:hi[]]) end return (NoTangent(), NoTangent(), dAs) end @@ -272,10 +278,11 @@ function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVe end function rrule(::typeof(reduce), ::typeof(hcat), As::AbstractVector{<:AbstractVector}) - axe = axes(As,1) - function reduce_hcat_pullback_1(dY) + project_As = ProjectTo.(As) + function reduce_hcat_pullback_1(ȳ) + dY = unthunk(ȳ) hi = Ref(0) - dAs = map(_ -> dY[:, hi[]+=1], axe) + dAs = map(project -> project(dY[:, hi[]+=1]), project_As) return (NoTangent(), NoTangent(), dAs) end return reduce(hcat, As), reduce_hcat_pullback_1 @@ -328,13 +335,15 @@ function rrule(::typeof(reduce), ::typeof(vcat), As::AbstractVector{<:AbstractVe Base.require_one_based_indexing(Y) ndimsY = Val(ndims(Y)) heights = map(A -> size(A,1), As) - function reduce_vcat_pullback(dY) + project_As = ProjectTo.(As) + function reduce_vcat_pullback(ȳ) + dY = unthunk(ȳ) hi = Ref(0) - dAs = map(heights) do z + dAs = map(project_As, heights) do project, z lo = hi[]+1 hi[] += z ind = ntuple(d -> d==1 ? (lo:hi[]) : (:), ndimsY) - dY[ind...] + project(dY[ind...]) end return (NoTangent(), NoTangent(), dAs) end @@ -398,12 +407,13 @@ function rrule(::typeof(hvcat), rows, values...) ndimsY = Val(ndims(Y)) sizes = map(_catsize, values) project_Vs = map(ProjectTo, values) - function hvcat_pullback(dY) + function hvcat_pullback(ȳ) + dY = unthunk(ȳ) prev = fill(0, 2) dXs = map(project_Vs, sizes) do project, sizeX ndimsX = length(sizeX) index = ntuple(ndimsY) do d - if d in (1, 2) + if d <= 2 d > ndimsX ? (prev[d]+1) : (prev[d]+1:prev[d]+sizeX[d]) else d > ndimsX ? 1 : (:) @@ -414,7 +424,10 @@ function rrule(::typeof(hvcat), rows, values...) prev[2] = 0 prev[1] += get(sizeX, 1, 1) end - project(dY[index...]) + InplaceableThunk( + dX -> dX .+= view(dY, index...), + @thunk project(@allowscalar dY[index...]) + ) end return (NoTangent(), NoTangent(), dXs...) end diff --git a/test/rulesets/Base/array.jl b/test/rulesets/Base/array.jl index 46a52a0a8..b94cf2600 100644 --- a/test/rulesets/Base/array.jl +++ b/test/rulesets/Base/array.jl @@ -199,7 +199,7 @@ end test_rrule(reduce, hcat, vecs) mix = AbstractVecOrMat[rand(4,2), rand(4)] # this is weird, but does hit the fast path - test_rrule(reduce, hcat, mix) + test_rrule(reduce, hcat, mix; check_inferred=false) adjs = vec([randn(2, 4), randn(1, 4), randn(3, 4)]') # not a Vector # test_rrule(reduce, hcat, adjs ⊢ map(m -> rand(size(m)), adjs)) @@ -240,7 +240,18 @@ end test_rrule(reduce, vcat, vecs) mix = AbstractVecOrMat[rand(4,1), rand(4)] - test_rrule(reduce, vcat, mix) + test_rrule(reduce, vcat, mix; check_inferred=false) +end + +@testset "reduce cat: cotangents are projected" begin + @testset "$red" for (red, As) in [ + (hcat, [rand(Float32, 3) for _ in 1:2]), + (hcat, AbstractVecOrMat[rand(Float32, 3, 2), rand(Float32, 3)]), + (vcat, AbstractVecOrMat[rand(Float32, 3, 1), rand(Float32, 3)]), + ] + Y, back = rrule(reduce, red, As) + @test map(typeof, back(rand(ComplexF64, size(Y)))[3]) == map(typeof, As) + end end @testset "cat" begin @@ -277,6 +288,33 @@ end test_rrule(hvcat, (1,2), rand(2)', (3.4, 5.6), 7.8; check_inferred=false) end +@testset "thunked cotangents are unthunked once" begin + # Unthunking recomputes the array, so the pullbacks must do so only once, not per input. + cases = [ + ("reduce vcat, vectors", rrule(reduce, vcat, [rand(3), rand(2), rand(4)])), + ("reduce vcat, matrices", rrule(reduce, vcat, [rand(2, 4), rand(1, 4), rand(3, 4)])), + ("reduce vcat, mixed", rrule(reduce, vcat, AbstractVecOrMat[rand(4, 1), rand(4)])), + ("reduce hcat, vectors", rrule(reduce, hcat, [rand(3) for _ in 1:4])), + ("reduce hcat, matrices", rrule(reduce, hcat, [rand(3, 2), rand(3, 1), rand(3, 3)])), + ("reduce hcat, mixed", rrule(reduce, hcat, AbstractVecOrMat[rand(4, 2), rand(4)])), + ("hvcat", rrule(hvcat, (2, 1), rand(), rand(1, 1), rand(2, 2))), + ("vect, uniform", rrule(Base.vect, (1.0,), (2.0,))), + ("vect, projected", rrule(Base.vect, 1.0, 2f0, 3)), + ("vect, any", rrule(Base.vect, :a, 1.0)), + ] + @testset "$name" for (name, (Y, back)) in cases + dY = rand(size(Y)...) + expected = map(unthunk, back(dY)) + n = Ref(0) + counted() = (n[] += 1; copy(dY)) + for dȲ in (@thunk(counted()), InplaceableThunk(dx -> dx .+= dY, @thunk(counted()))) + n[] = 0 + @test map(unthunk, back(dȲ)) == expected + @test n[] == 1 + end + end +end + @testset "reverse" begin @testset "Tuple" begin test_frule(reverse, Tuple(rand(10)))