From f4fa2e70291c7c565cfe237ee3d5beb93123c08d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 14:34:47 -0400 Subject: [PATCH] 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]