Skip inferring the trailing-dims reduction in tuple unbroadcast for vectors

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"

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)

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]