[lkdvos]

unbroadcast(x::Tuple, dx) only needs to reduce over trailing dimensions when a tuple was broadcast against something with more than one dimension:

val = if N == length(dx)
    dx
else
    sum(dx; dims=2:ndims(dx))
end

For a vector dx that branch is dead. A tuple of length N ≥ 2 (Tuple{Any} has its own method) fills the first dimension, so length(dx) == N always holds, and sum(dx; dims=2:1) would be a no-op anyway. Inference can't see this from the runtime N == length(dx) check, though, so it still compiles sum(dx; dims=...) for the eltype of dx. In the generic broadcast path (split_bc_pullbacks) that eltype is often abstract.

This PR checks dx isa AbstractVector first, so the branch is pruned statically. Runtime behaviour doesn't change.

Minimal reproducer

using ChainRules
@time ChainRules.unbroadcast((1.0, 2.0), [3.0, 4.0])

First call, fresh session, Julia 1.13.3:

dx	before	after
Vector{Float64}	0.155 s	0.074 s
Vector{Any}	0.175 s	0.104 s
Vector{AbstractVector{Float64}}	0.188 s	0.119 s

Downstream impact

The dead reduction becomes a real problem when inferring + on the abstract eltype is expensive. In QuantumKitHub/BlockTensorKit.jl#86, Zygote.gradient of x -> (a, b = adjoint.(x); real(dot(a, a) + dot(b, b))) on a tuple of TensorKit tensors goes like this:

• Zygote passes a Tangent{Any, Tuple{...}} cotangent, so back_generic collects the per-element tangents into a Vector{<:AdjointTensorMap{...}} with an abstract eltype.
• Inferring the never-executed sum over that vector then takes several minutes when BlockTensorKit is loaded.

With this change the first gradient takes 7.1 s (>150 s before), with no changes to TensorKit or BlockTensorKit.

Tests for the vector, abstract-eltype and matrix (still reduced) cases are added next to the existing unbroadcast with NTuple tests; rulesets/Base/broadcast passes locally. Version bumped to 1.74.9.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.63%. Comparing base (7fb6cb1) to head (f4fa2e7).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #857   +/-   ##
=======================================
  Coverage   96.63%   96.63%           
=======================================
  Files          28       28           
  Lines        3596     3596           
=======================================
  Hits         3475     3475           
  Misses        121      121           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lkdvos]

As a manual follow-up here (TLDR):

• length(dx) == N on a AbstractVector is always true, but the compiler cannot prove it
• The sum(dx; dims=2:...) branch takes a long time to infer on some inputs of abstract types
• Zygote produces Tangent{Any, Tuple{...}} cotangents, which result in per-element tangents of a tuple broadcast to be collected into a Vector (hence the AbstractVector case being important)

[devmotion]

AFAICT the change is correct, but I don't think it addresses the actual problem in the BlockTensorKit example. There the broadcast is tuple-to-tuple, so the cotangent should be a tuple and not a Vector. It only ends up as a Vector because Tangent doesn't define broadcastable and hence broadcasting falls back to collect. I opened JuliaDiff/ChainRulesCore.jl#724 which makes Tangents broadcast like their primal: tuple tangents like their backing tuple, all other tangents as scalars. With that PR unbroadcast receives a Tuple in the example and the vector branch is never hit. It also fixes broadcasts that return a struct, e.g. f.(Ref(x)), which currently fail with Need an adjoint for constructor ....

Vector cotangents still show up if a tuple is broadcast with a vector, e.g. (a, b) .* [1, 2], so avoiding the reduction is still useful in this case. But IMO it would be cleaner not to special-case vectors and instead use that a tuple only ever spans the first dimension: whether we have to reduce depends only on ndims(dx), not on length(dx). See my comment below.

Unrelated to this PR, but I noticed that Tuple{Vararg{Any,N}}(val) silently drops elements if the cotangent is too long (also on main):

julia> ChainRules.unbroadcast((1.0, 2.0), [3.0, 4.0, 5.0])
Tangent{Tuple{Float64, Float64}}(3.0, 4.0)

This can't happen with correct rules, but it hides bugs in other rules. Maybe we should check the length and throw a DimensionMismatch?

[devmotion]

What about dispatching on ndims instead?

function unbroadcast(x::T, dx_raw) where {T<:Tuple{Vararg{Any,N}}} where {N}
    val = _sum_trailing_dims(unthunk(dx_raw))
    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, so only trailing dimensions need summing.
_sum_trailing_dims(dx::AbstractArray) = ndims(dx) == 1 ? dx : sum(dx; dims=2:ndims(dx))
_sum_trailing_dims(dx) = dx  # e.g. `Tuple`, `Tangent{<:Any,<:Tuple}`

This avoids the reduction for vectors, tuples and Tangents but still sums over trailing dimensions for matrices. I checked locally that it gives the same results as main (vectors, Vector{Any}, tuples, Tangents, N×k, N×1 and 3d arrays, thunks, zero cotangents), and the first call is as fast as with this PR.

[devmotion]

The test doesn't check that the reduction is not compiled, it only compares values. Can you remove the comment, or test it?

[lkdvos]

I'd definitely vote in favor of the dimensionmismatch. For the ndims I'm a bit more on the fence, it seems to me that you can probably broadcast a lot of things that don't necessarily have ndims defined, which is why I kind of hesitated to use that here as well. I'm not sure if the base interface for detecting iteration shape is correct either, since that doensnt necessarily coincide with broadcast either (I think? It's always hard to know what contracts are meant with the Base interfaces)

[devmotion]

A quick update, since part of my review above is out of date. JuliaDiff/ChainRulesCore.jl#724 no longer makes Tangents broadcast like their primal. No broadcastable definition is correct for all Tangents, so broadcasting over structural tangents now throws an ArgumentError. Unwrapping the cotangent happens in the broadcast rules instead, so I reopened #858.

This PR is still useful. Vector cotangents still show up when a tuple is broadcast against a vector, e.g. (a, b) .* [1, 2].

On ndims: my suggestion only calls ndims on AbstractArrays, which always define it. Anything else, such as a Tuple or a Tangent, goes to the fallback and doesn't need ndims. A tuple only ever spans the first dimension of a broadcast, so whether to reduce depends only on the number of dimensions of dx.

Since you're in favour of the DimensionMismatch, I'd add it in this PR too.

What's left: the test comment "reduction not compiled" (it should either be removed or tested), and a rebase on main.