diff --git a/Project.toml b/Project.toml index faf9a412c..6f9e69e11 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.15" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Base/broadcast.jl b/src/rulesets/Base/broadcast.jl index ae389b84f..debf46a67 100644 --- a/src/rulesets/Base/broadcast.jl +++ b/src/rulesets/Base/broadcast.jl @@ -320,6 +320,7 @@ rrule(::typeof(broadcasted), ::typeof(complex), x::Number) = rrule(complex, x) | function unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx_raw) dx = unthunk(dx_raw) + dx isa AbstractZero && return dx N = _ndims(dx) if length(x) == length(dx) ProjectTo(x)(dx) # handles trivial reshapes, offsets, structured matrices, row vectors @@ -328,22 +329,24 @@ function unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx_raw) ProjectTo(x)(sum(dx; dims)) end end -unbroadcast(x::Base.AbstractArrayOrBroadcasted, dx::AbstractZero) = dx _ndims(x) = ndims(x) _ndims(::Tuple) = 1 function unbroadcast(x::T, dx_raw) where {T<:Tuple{Vararg{Any,N}}} where {N} - dx = unthunk(dx_raw) - val = if N == length(dx) - dx - else - sum(dx; dims=2:ndims(dx)) - end + val = _sum_trailing_dims(unthunk(dx_raw)) + val isa AbstractZero && return val + n = length(val) + n == N || + throw(DimensionMismatch(lazy"cotangent of length $n for a tuple of length $N")) eltype(val) <: AbstractZero && return NoTangent() return ProjectTo(x)(Tuple{Vararg{Any,N}}(val)) # Tangent end -unbroadcast(x::Tuple, dx::AbstractZero) = dx + +# A tuple only spans the first dimension of a broadcast +_sum_trailing_dims(dx::AbstractVector) = dx +_sum_trailing_dims(dx::AbstractArray) = sum(dx; dims=2:ndims(dx)) +_sum_trailing_dims(dx) = dx # e.g. `Tuple`, `Tangent{<:Any,<:Tuple}` # zero(::Tangent) is some Zero, which means sum(dx; dims) fails unless you do this: function Base.reducedim_init( @@ -381,7 +384,6 @@ unbroadcast(x::Base.RefValue, dx::AbstractZero) = dx unbroadcast(::Bool, dx) = NoTangent() unbroadcast(::AbstractArray{Bool}, dx) = NoTangent() -unbroadcast(::AbstractArray{Bool}, dx::AbstractZero) = dx # ambiguity unbroadcast(::Val, dx) = NoTangent() function unbroadcast(f::Function, df) diff --git a/test/rulesets/Base/broadcast.jl b/test/rulesets/Base/broadcast.jl index ea9bac8d0..9f70dae2b 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -201,6 +201,31 @@ BT1 = Broadcast.BroadcastStyle(Tuple) @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) end + @testset "unbroadcast tuple" begin + x = (1.0, 2.0) + dx = Tangent{typeof(x)}(3.0, 4.0) + @test @inferred(ChainRules.unbroadcast(x, [3.0, 4.0])) == dx + @test ChainRules.unbroadcast(x, Any[3.0, 4.0]) == dx # not inferable: the eltype may be a zero + @test @inferred(ChainRules.unbroadcast(x, [3.0; 4.0;;])) == dx + @test @inferred(ChainRules.unbroadcast(x, [1.0 2.0; 1.5 2.5])) == dx + @test @inferred( + ChainRules.unbroadcast(x, reshape([1.0, 2.0, 2.0, 2.0], 2, 1, 2)) + ) == dx + @test @inferred(ChainRules.unbroadcast(x, (3.0, 4.0))) == dx + @test @inferred(ChainRules.unbroadcast(x, Tangent{Any}(3.0, 4.0))) == dx + @test @inferred(ChainRules.unbroadcast(x, @thunk([3.0, 4.0]))) == dx + @test @inferred(ChainRules.unbroadcast(x, ZeroTangent())) == ZeroTangent() + @test @inferred(ChainRules.unbroadcast(x, @thunk(ZeroTangent()))) == + ZeroTangent() + @test @inferred(ChainRules.unbroadcast([1.0, 2.0], @thunk(ZeroTangent()))) == + ZeroTangent() + @test @inferred(ChainRules.unbroadcast(x, [ZeroTangent(), ZeroTangent()])) == + NoTangent() + @test_throws DimensionMismatch ChainRules.unbroadcast(x, [3.0, 4.0, 5.0]) + @test_throws DimensionMismatch ChainRules.unbroadcast( + x, [1.0 2.0; 1.5 2.5; 0.0 0.0] + ) + end @testset "unbroadcast with Matrix{Tangent}" begin # https://github.com/JuliaDiff/ChainRules.jl/issues/708 x = Base.Fix1.(*, 1:3.0) dx1 = [Tangent{Base.Fix1}(; x = i/2) for i in 1:3, _ in 1:1]