diff --git a/Project.toml b/Project.toml index faf9a412c..1d8c5e7b6 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.13" +version = "1.74.14" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Base/broadcast.jl b/src/rulesets/Base/broadcast.jl index ae389b84f..16ba6466e 100644 --- a/src/rulesets/Base/broadcast.jl +++ b/src/rulesets/Base/broadcast.jl @@ -2,6 +2,12 @@ using Base.Broadcast: Broadcast, broadcasted, Broadcasted, BroadcastStyle const RCR = RuleConfig{>:HasReverseMode} const TRI_NO = (NoTangent(), NoTangent(), NoTangent()) +# The cotangent of a tuple-valued broadcast can be a `Tangent` (`Tangent{Any}` from Zygote). +# Broadcast over its backing tuple, since broadcasting over a `Tangent` either loses the element +# types or throws. +_bc_cotangent(dys::Tangent{<:Any,<:Tuple}) = ChainRulesCore.backing(dys) +_bc_cotangent(dys) = dys + function rrule(::typeof(copy), bc::Broadcasted) uncopy(Δ) = (NoTangent(), Δ) return copy(bc), uncopy @@ -59,7 +65,7 @@ function split_bc_derivatives(f::F, arg) where {F} @debug("split broadcasting derivative", f) ys = f.(arg) function bc_one_back(dys) # For f.(x) we do not need StructArrays / unzip at all - delta = broadcast(unthunk(dys), ys, arg) do dy, y, a + delta = broadcast(_bc_cotangent(unthunk(dys)), ys, arg) do dy, y, a das = only(derivatives_given_output(y, f, a)) dy * conj(only(das)) # possibly this * should be made nan-safe. end @@ -72,7 +78,7 @@ function split_bc_derivatives(f::F, args::Vararg{Any,N}) where {F,N} @debug("split broadcasting derivatives", f, N) ys = f.(args...) function bc_many_back(dys) - deltas = unzip_broadcast(unthunk(dys), ys, args...) do dy, y, as... + deltas = unzip_broadcast(_bc_cotangent(unthunk(dys)), ys, args...) do dy, y, as... das = only(derivatives_given_output(y, f, as...)) map(da -> dy * conj(da), das) # possibly this * should be made nan-safe. end @@ -110,7 +116,7 @@ function split_bc_inner(frule_fun::R, cfg::RuleConfig, f::F, arg) where {R,F} frule_fun(cfg, (NoTangent(), one(a)), f, a) end function back_forwards(dys) - delta = broadcast(ydots, unthunk(dys), arg) do ydot, dy, a + delta = broadcast(ydots, _bc_cotangent(unthunk(dys)), arg) do ydot, dy, a ProjectTo(a)(conj(ydot) * dy) # possibly this * should be made nan-safe. end return (TRI_NO..., ProjectTo(arg)(delta)) @@ -129,7 +135,7 @@ function split_bc_pullbacks(cfg::RCR, f::F, args::Vararg{Any,N}) where {F,N} rrule_via_ad(cfg, f, a...) end function back_generic(dys) - deltas = unzip_broadcast(backs, dys) do back, dy # (could be map, sizes match) + deltas = unzip_broadcast(backs, _bc_cotangent(dys)) do back, dy # (could be map, sizes match) map(unthunk, back(dy)) end dargs = map(unbroadcast, args, Base.tail(deltas)) @@ -162,7 +168,7 @@ const NumericOrBroadcast = Union{Number, AbstractArray{<:Number}, NTuple{<:Any,N function rrule(::typeof(broadcasted), ::typeof(+), xs::NumericOrBroadcast...) @debug("broadcasting: plus", length(xs)) function bc_plus_back(dy_raw) - dy = unthunk(dy_raw) + dy = _bc_cotangent(unthunk(dy_raw)) return (NoTangent(), NoTangent(), map(x -> unbroadcast(x, dy), xs)...) # no copies, this may return dx2 === dx3 end return broadcasted(+, xs...), bc_plus_back @@ -171,7 +177,7 @@ end function rrule(::typeof(broadcasted), ::typeof(-), x::NumericOrBroadcast, y::NumericOrBroadcast) @debug("broadcasting: minus 2") function bc_minus_back(dz_raw) - dz = unthunk(dz_raw) + dz = _bc_cotangent(unthunk(dz_raw)) return (NoTangent(), NoTangent(), @thunk(unbroadcast(x, dz)), @thunk(-unbroadcast(y, dz))) end return broadcasted(-, x, y), bc_minus_back @@ -179,20 +185,21 @@ end function rrule(::typeof(broadcasted), ::typeof(-), x::NumericOrBroadcast) @debug("broadcasting: minus 1") - bc_minus_back(dy) = (NoTangent(), NoTangent(), @thunk -unthunk(dy)) + bc_minus_back(dy) = + (NoTangent(), NoTangent(), unbroadcast(x, .-_bc_cotangent(unthunk(dy)))) return broadcasted(-, x), bc_minus_back end function rrule(::typeof(broadcasted), ::typeof(*), x::NumericOrBroadcast, y::NumericOrBroadcast) @debug("broadcasting: times") function bc_times_back(Δraw) - Δ = unthunk(Δraw) + Δ = _bc_cotangent(unthunk(Δraw)) return (NoTangent(), NoTangent(), _back_star(x, y, Δ), _back_star(y, x, Δ)) end return broadcasted(*, x, y), bc_times_back end _back_star(x, y, Δ) = @thunk unbroadcast(x, Δ .* conj.(y)) # this case probably isn't better than generic -_back_star(x::Number, y, Δ) = @thunk LinearAlgebra.dot(y, Δ) # ... but this is why the rule exists +_back_star(x::Number, y, Δ) = @thunk unbroadcast(x, LinearAlgebra.dot(y, Δ)) # ... but this is why the rule exists _back_star(x::Bool, y, Δ) = NoTangent() _back_star(x::Complex{Bool}, y, Δ) = NoTangent() # e.g. for fun.(im.*x) @@ -212,7 +219,7 @@ end function rrule(::typeof(broadcasted), ::typeof(Base.literal_pow), ::typeof(^), x::NumericOrBroadcast, ::Val{2}) @debug("broadcasting: square") function bc_square_back(dy_raw) - dx = @thunk ProjectTo(x)(2 .* unthunk(dy_raw) .* conj.(x)) + dx = unbroadcast(x, 2 .* _bc_cotangent(unthunk(dy_raw)) .* conj.(x)) return (NoTangent(), NoTangent(), NoTangent(), dx, NoTangent()) end return broadcasted(Base.literal_pow, ^, x, Val(2)), bc_square_back @@ -223,10 +230,14 @@ function rrule(::typeof(broadcasted), ::typeof(/), x::NumericOrBroadcast, y::Num # z = broadcast(/, x, y) z = broadcasted(/, x, y) function bc_divide_back(dz_raw) - dz = unthunk(dz_raw) + dz = _bc_cotangent(unthunk(dz_raw)) dx = @thunk unbroadcast(x, dz ./ conj.(y)) # dy = @thunk -LinearAlgebra.dot(z, dz) / conj(y) # the reason to be eager is to allow dot here - dy = @thunk -sum(Broadcast.instantiate(broadcasted(*, broadcasted(conj, z), dz))) / conj(y) # complete sum is fast + # complete sum is fast + dy = @thunk unbroadcast( + y, + -sum(Broadcast.instantiate(broadcasted(*, broadcasted(conj, z), dz))) / conj(y), + ) return (NoTangent(), NoTangent(), dx, dy) end return z, bc_divide_back @@ -257,14 +268,16 @@ rrule(::typeof(broadcasted), ::typeof(identity), x::Number) = rrule(identity, x) function rrule(::typeof(broadcasted), ::Type{T}, x::NumericOrBroadcast) where {T<:Number} @debug("broadcasting: type", T) - bc_type_back(dz) = (NoTangent(), NoTangent(), @thunk(unbroadcast(x, unthunk(dz)))) + bc_type_back(dz) = + (NoTangent(), NoTangent(), unbroadcast(x, _bc_cotangent(unthunk(dz)))) return broadcasted(T, x), bc_type_back end rrule(::typeof(broadcasted), ::Type{T}, x::Number) where {T<:Number} = rrule(T, x) |> _prepend_zero function rrule(::typeof(broadcasted), ::typeof(float), x::NumericOrBroadcast) @debug("broadcasting: float") - bc_float_back(dz) = (NoTangent(), NoTangent(), @thunk(unbroadcast(x, unthunk(dz)))) + bc_float_back(dz) = + (NoTangent(), NoTangent(), unbroadcast(x, _bc_cotangent(unthunk(dz)))) return broadcasted(float, x), bc_float_back end rrule(::typeof(broadcasted), ::typeof(float), x::Number) = rrule(float, x) |> _prepend_zero @@ -274,7 +287,11 @@ rrule(::typeof(broadcasted), ::typeof(float), x::Number) = rrule(float, x) |> _p for conj in [:conj, :adjoint] # identical as we know eltype <: Number @eval begin function rrule(::typeof(broadcasted), ::typeof($conj), x::NumericOrBroadcast) - bc_conj_back(dx) = (NoTangent(), NoTangent(), conj(unthunk(dx))) + bc_conj_back(dx) = ( + NoTangent(), + NoTangent(), + unbroadcast(x, conj.(_bc_cotangent(unthunk(dx)))), + ) return broadcasted($conj, x), bc_conj_back end rrule(::typeof(broadcasted), ::typeof($conj), x::Number) = rrule($conj, x) |> _prepend_zero @@ -286,15 +303,19 @@ end function rrule(::typeof(broadcasted), ::typeof(real), x::NumericOrBroadcast) @debug("broadcasting: real") - bc_real_back(dz) = (NoTangent(), NoTangent(), @thunk(real(unthunk(dz)))) + bc_real_back(dz) = + (NoTangent(), NoTangent(), unbroadcast(x, _real(_bc_cotangent(unthunk(dz))))) return broadcasted(real, x), bc_real_back end +_real(dz) = real(dz) # doesn't copy a real array +_real(dz::Tuple) = real.(dz) rrule(::typeof(broadcasted), ::typeof(real), x::Number) = rrule(real, x) |> _prepend_zero rrule(::typeof(broadcasted), ::typeof(real), x::AbstractArray{<:Real}) = rrule(identity, x) |> _prepend_zero function rrule(::typeof(broadcasted), ::typeof(imag), x::NumericOrBroadcast) @debug("broadcasting: imag") - bc_imag_back(dz) = (NoTangent(), NoTangent(), @thunk(im .* real.(unthunk(dz)))) + bc_imag_back(dz) = + (NoTangent(), NoTangent(), unbroadcast(x, im .* real.(_bc_cotangent(unthunk(dz))))) return broadcasted(imag, x), bc_imag_back end rrule(::typeof(broadcasted), ::typeof(imag), x::Number) = rrule(imag, x) |> _prepend_zero @@ -306,7 +327,8 @@ end function rrule(::typeof(broadcasted), ::typeof(complex), x::NumericOrBroadcast) @debug("broadcasting: complex") - bc_complex_back(dz) = (NoTangent(), NoTangent(), @thunk(unbroadcast(x, unthunk(dz)))) + bc_complex_back(dz) = + (NoTangent(), NoTangent(), unbroadcast(x, _bc_cotangent(unthunk(dz)))) return broadcasted(complex, x), bc_complex_back end rrule(::typeof(broadcasted), ::typeof(complex), x::Number) = rrule(complex, x) |> _prepend_zero diff --git a/test/rulesets/Base/broadcast.jl b/test/rulesets/Base/broadcast.jl index ea9bac8d0..07bd1710f 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -92,13 +92,17 @@ BT1 = Broadcast.BroadcastStyle(Tuple) @gpu test_rrule(copy∘broadcasted, -, rand(3)) test_rrule(copy∘broadcasted, -, Tuple(rand(3))) - @gpu test_rrule(copy∘broadcasted, *, rand(3), rand(3)) - @gpu test_rrule(copy∘broadcasted, *, rand(3), rand()) - @gpu test_rrule(copy∘broadcasted, *, rand(), rand(3)) + @gpu test_rrule(copy ∘ broadcasted, *, rand(3), rand(3)) + @gpu test_rrule(copy ∘ broadcasted, *, rand(3), rand()) + @gpu test_rrule(copy ∘ broadcasted, *, rand(), rand(3)) + test_rrule(copy ∘ broadcasted, *, Tuple(rand(3)), rand()) + test_rrule(copy ∘ broadcasted, *, rand(), Tuple(rand(3))) + test_rrule(copy ∘ broadcasted, *, Tuple(rand(3)), Tuple(rand(3))) test_rrule(copy∘broadcasted, *, rand(3) .+ im, rand(3) .+ 2im) test_rrule(copy∘broadcasted, *, rand(3) .+ im, rand() + 3im) test_rrule(copy∘broadcasted, *, rand() + im, rand(3) .+ 4im) + test_rrule(copy ∘ broadcasted, *, rand(), rand(3) .+ im) # projects the cotangent of the real scalar @test_skip test_rrule(copy∘broadcasted, *, im, rand(3)) # MethodError: no method matching randn(::Random._GLOBAL_RNG, ::Type{Complex{Bool}}) @test_skip test_rrule(copy∘broadcasted, *, rand(3), im) # MethodError: no method matching randn(::Random._GLOBAL_RNG, ::Type{Complex{Bool}}) @@ -119,11 +123,14 @@ BT1 = Broadcast.BroadcastStyle(Tuple) # GPU error from dot(x::JLArray{Float32, 1}, y::JLArray{ComplexF32, 2}) end - @gpu test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3), Val(2)) - @gpu test_rrule(copy∘broadcasted, Base.literal_pow, ^, rand(3) .+ im, Val(2)) + @gpu test_rrule(copy ∘ broadcasted, Base.literal_pow, ^, rand(3), Val(2)) + @gpu test_rrule(copy ∘ broadcasted, Base.literal_pow, ^, rand(3) .+ im, Val(2)) + test_rrule(copy ∘ broadcasted, Base.literal_pow, ^, Tuple(rand(3)), Val(2)) - @gpu test_rrule(copy∘broadcasted, /, rand(3), rand()) - @gpu test_rrule(copy∘broadcasted, /, rand(3) .+ im, rand() + 3im) + @gpu test_rrule(copy ∘ broadcasted, /, rand(3), rand()) + @gpu test_rrule(copy ∘ broadcasted, /, rand(3) .+ im, rand() + 3im) + test_rrule(copy ∘ broadcasted, /, Tuple(rand(3)), rand()) + test_rrule(copy ∘ broadcasted, /, rand(3) .+ im, rand()) # projects the cotangent of the real scalar end @testset "identity etc" begin test_rrule(copy∘broadcasted, identity, rand(3)) @@ -134,16 +141,22 @@ BT1 = Broadcast.BroadcastStyle(Tuple) test_rrule(copy∘broadcasted, float, rand(3)) end @testset "complex" begin - test_rrule(copy∘broadcasted, conj, rand(3)) - test_rrule(copy∘broadcasted, conj, rand(3) .+ im) - test_rrule(copy∘broadcasted, adjoint, rand(3)) - test_rrule(copy∘broadcasted, adjoint, rand(3) .+ im) + test_rrule(copy ∘ broadcasted, conj, rand(3)) + test_rrule(copy ∘ broadcasted, conj, rand(3) .+ im) + test_rrule(copy ∘ broadcasted, conj, Tuple(rand(3))) + test_rrule(copy ∘ broadcasted, conj, Tuple(rand(3) .+ im .* rand.())) + test_rrule(copy ∘ broadcasted, adjoint, rand(3)) + test_rrule(copy ∘ broadcasted, adjoint, rand(3) .+ im) - test_rrule(copy∘broadcasted, real, rand(3)) - test_rrule(copy∘broadcasted, real, rand(3) .+ im) + test_rrule(copy ∘ broadcasted, real, rand(3)) + test_rrule(copy ∘ broadcasted, real, rand(3) .+ im) + test_rrule(copy ∘ broadcasted, real, Tuple(rand(3))) + test_rrule(copy ∘ broadcasted, real, Tuple(rand(3) .+ im .* rand.())) - test_rrule(copy∘broadcasted, imag, rand(3)) - test_rrule(copy∘broadcasted, imag, rand(3) .+ im .* rand.()) + test_rrule(copy ∘ broadcasted, imag, rand(3)) + test_rrule(copy ∘ broadcasted, imag, rand(3) .+ im .* rand.()) + test_rrule(copy ∘ broadcasted, imag, Tuple(rand(3))) + test_rrule(copy ∘ broadcasted, imag, Tuple(rand(3) .+ im .* rand.())) test_rrule(copy∘broadcasted, complex, rand(3)) end @@ -197,6 +210,70 @@ BT1 = Broadcast.BroadcastStyle(Tuple) @test y3 == () @test bk3(Tangent{Tuple{}}()) == (NoTangent(), NoTangent(), NoTangent(), ZeroTangent(), ZeroTangent()) end + @testset "tuple broadcast with Tangent{Any} cotangent" begin + x = (1.0, 2.0f0) + y2, bk2 = rrule(CFG, copy ∘ broadcasted, BT1, log, x) # split 2 + dx2 = @inferred bk2(Tangent{Any}(1.0, 1.0f0)) + @test dx2[4] isa Tangent{typeof(x),typeof(x)} + @test all(ChainRulesCore.backing(dx2[4]) .≈ (1.0, 0.5)) + + y4, bk4 = rrule(CFG, copy ∘ broadcasted, BT1, sin ∘ cos, x) # split 4 + dx4 = bk4(Tangent{Any}(1.0, 1.0f0)) + @test dx4[4] isa Tangent{typeof(x),typeof(x)} + @test all( + ChainRulesCore.backing(dx4[4]) .≈ + (-cos(cos(1.0)) * sin(1.0), -cos(cos(2.0f0)) * sin(2.0f0)), + ) + + @test bk2(@thunk(Tangent{Any}(1.0, 1.0f0))) == dx2 + @test bk4(@thunk(Tangent{Any}(1.0, 1.0f0))) == dx4 + + _, bkm = rrule(CFG, copy ∘ broadcasted, BT1, atan, x, x) # split 2, several arguments + dxm = @inferred bkm(Tangent{Any}(1.0, 1.0f0)) + @test dxm[4] isa Tangent{typeof(x),typeof(x)} + @test dxm[5] isa Tangent{typeof(x),typeof(x)} + + x3 = (1.0, 2.0) # split 3 needs a concrete eltype + _, bk3 = rrule(CFG, copy ∘ broadcasted, BT1, flog, x3) + dx3 = bk3(Tangent{Any}(1.0, 1.0)) + @test dx3[4] isa Tangent{typeof(x3),typeof(x3)} + @test all(ChainRulesCore.backing(dx3[4]) .≈ (1.0, 0.5)) + + # Fused rules + dy = Tangent{Any}(1.0, 1.0f0) + _, _, dxt, dyt = map(unthunk, rrule(CFG, copy ∘ broadcasted, *, x, 2.0)[2](dy)) + @test dxt == Tangent{typeof(x)}(2.0, 2.0f0) + @test dyt ≈ 3.0 + _, _, _, dxs, _ = map( + unthunk, + rrule(CFG, copy ∘ broadcasted, Base.literal_pow, ^, x, Val(2))[2](dy), + ) + @test dxs == Tangent{typeof(x)}(2.0, 4.0f0) + _, _, dxd, dyd = map(unthunk, rrule(CFG, copy ∘ broadcasted, /, x, 2.0)[2](dy)) + @test dxd == Tangent{typeof(x)}(0.5, 0.5f0) + @test dyd ≈ -0.75 + xc = (1.0 + 2im, 3.0 - 1im) + @test rrule(CFG, copy ∘ broadcasted, real, xc)[2](dy)[3] == + Tangent{typeof(xc)}(1.0 + 0im, 1.0 + 0im) + @test rrule(CFG, copy ∘ broadcasted, imag, xc)[2](dy)[3] == + Tangent{typeof(xc)}(1.0im, 1.0im) + dyc = Tangent{Any}(1.0 + 1im, 2.0 - 1im) + @test rrule(CFG, copy ∘ broadcasted, conj, xc)[2](dyc)[3] == + Tangent{typeof(xc)}(1.0 - 1im, 2.0 + 1im) + # Zygote passes the cotangent of the outer rule in e.g. `real.(conj.(x))` on as a `Tuple` + @test rrule(CFG, copy ∘ broadcasted, conj, xc)[2]((1.0, 1.0))[3] == + Tangent{typeof(xc)}(1.0 + 0im, 1.0 + 0im) + @test rrule(CFG, copy ∘ broadcasted, -, x)[2]((1.0, 1.0f0))[3] == + Tangent{typeof(x)}(-1.0, -1.0f0) + + # Lazy arguments, e.g. from `float.(real.(x))` + xl = broadcasted(real, x) + @test unthunk(rrule(broadcasted, +, xl, 1.0)[2](dy)[3]) == (1.0, 1.0f0) + @test unthunk(rrule(broadcasted, -, xl, 1.0)[2](dy)[3]) == (1.0, 1.0f0) + @test rrule(broadcasted, float, xl)[2](dy)[3] == (1.0, 1.0f0) + @test rrule(broadcasted, Float32, xl)[2](dy)[3] == (1.0, 1.0f0) + @test rrule(broadcasted, complex, xl)[2](dy)[3] == (1.0, 1.0f0) + end @testset "unbroadcast with NTuple" begin # https://github.com/JuliaDiff/ChainRules.jl/pull/661 @test ChainRules.unbroadcast((1, 2, [3]), [4, 5, [6]]) isa Tangent # earlier, NTuple demanded same type @test ChainRules.unbroadcast(broadcasted(-, (1, 2), 3), (4, 5)) == (4, 5) # earlier, called ndims(::Tuple)