From f4fa2e70291c7c565cfe237ee3d5beb93123c08d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 14:34:47 -0400 Subject: [PATCH 1/3] Skip inferring the trailing-dims reduction in tuple `unbroadcast` for vectors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A tuple of length N ≥ 2 fills the first dimension of a broadcast, so a vector cotangent always has length N and the `sum(dx; dims=2:ndims(dx))` branch is dead. The runtime check `N == length(dx)` cannot be resolved by inference, which therefore still compiles that reduction for the (often abstract) eltype of `dx`. Checking `dx isa AbstractVector` first lets the branch be pruned statically. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- src/rulesets/Base/broadcast.jl | 4 +++- test/rulesets/Base/broadcast.jl | 6 ++++++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 2a59656eb..a18c5c4b0 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.8" +version = "1.74.9" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Base/broadcast.jl b/src/rulesets/Base/broadcast.jl index ae389b84f..9e58ed834 100644 --- a/src/rulesets/Base/broadcast.jl +++ b/src/rulesets/Base/broadcast.jl @@ -335,7 +335,9 @@ _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) + # A tuple of length N ≥ 2 fills the first dimension, so a vector `dx` already has length N. + # Checking this statically keeps inference from compiling the (dead) reduction below. + val = if dx isa AbstractVector || N == length(dx) dx else sum(dx; dims=2:ndims(dx)) diff --git a/test/rulesets/Base/broadcast.jl b/test/rulesets/Base/broadcast.jl index ea9bac8d0..6574a2329 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -201,6 +201,12 @@ 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 with vector cotangent" begin + dx = Tangent{Tuple{Float64, Float64}}(3.0, 4.0) + @test ChainRules.unbroadcast((1.0, 2.0), [3.0, 4.0]) == dx + @test ChainRules.unbroadcast((1.0, 2.0), Any[3.0, 4.0]) == dx # abstract eltype, reduction not compiled + @test ChainRules.unbroadcast((1.0, 2.0), [1.0 2.0; 1.5 2.5]) == dx # trailing dimensions still summed + 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] From 6586e3afbd5873aade04e63ba9f3d0ad19a7117f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 12 Nov 2026 01:16:55 +0100 Subject: [PATCH 2/3] Dispatch on the cotangent type in tuple `unbroadcast` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Sum over trailing dimensions in `_sum_trailing_dims`, which dispatches on the cotangent type: vectors, tuples and `Tangent`s are returned as is, so inference never sees the reduction for them. - Throw a `DimensionMismatch` for a cotangent of the wrong length, which was silently truncated before. - Return thunked zero cotangents early, in the tuple and the array method; they failed with `length(::ZeroTangent)` and `ndims(::ZeroTangent)`. This makes the `::AbstractZero` methods for tuples and arrays (and the `AbstractArray{Bool}` disambiguation) redundant. - Test N×1, 3-d, `Tuple`, `Tangent`, thunked and zero cotangents and the length mismatch, with `@inferred` where the result type is inferable. Co-Authored-By: Claude Opus 5.5 --- src/rulesets/Base/broadcast.jl | 21 ++++++++++----------- test/rulesets/Base/broadcast.jl | 29 ++++++++++++++++++++++++----- 2 files changed, 34 insertions(+), 16 deletions(-) diff --git a/src/rulesets/Base/broadcast.jl b/src/rulesets/Base/broadcast.jl index 9e58ed834..e764016ce 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,24 +329,23 @@ 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) - # A tuple of length N ≥ 2 fills the first dimension, so a vector `dx` already has length N. - # Checking this statically keeps inference from compiling the (dead) reduction below. - val = if dx isa AbstractVector || 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( @@ -383,7 +383,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 6574a2329..9f70dae2b 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -201,11 +201,30 @@ 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 with vector cotangent" begin - dx = Tangent{Tuple{Float64, Float64}}(3.0, 4.0) - @test ChainRules.unbroadcast((1.0, 2.0), [3.0, 4.0]) == dx - @test ChainRules.unbroadcast((1.0, 2.0), Any[3.0, 4.0]) == dx # abstract eltype, reduction not compiled - @test ChainRules.unbroadcast((1.0, 2.0), [1.0 2.0; 1.5 2.5]) == dx # trailing dimensions still summed + @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) From 98fee5744ad9425d451113c2a73a1a5e4b95a076 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 12 Nov 2026 01:30:00 +0100 Subject: [PATCH 3/3] Fix format --- src/rulesets/Base/broadcast.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/rulesets/Base/broadcast.jl b/src/rulesets/Base/broadcast.jl index e764016ce..debf46a67 100644 --- a/src/rulesets/Base/broadcast.jl +++ b/src/rulesets/Base/broadcast.jl @@ -337,7 +337,8 @@ function unbroadcast(x::T, dx_raw) where {T<:Tuple{Vararg{Any,N}}} where {N} 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")) + 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