Repository navigation
Conversation
… vectors 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 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #857 +/- ##
==========================================
+ Coverage 96.65% 96.73% +0.08%
==========================================
Files 28 28
Lines 3615 3616 +1
==========================================
+ Hits 3494 3498 +4
+ Misses 121 118 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
As a manual follow-up here (TLDR):
|
devmotion
left a comment
There was a problem hiding this comment.
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?
| 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) |
There was a problem hiding this comment.
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.
| @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 |
There was a problem hiding this comment.
The test doesn't check that the reduction is not compiled, it only compares values. Can you remove the comment, or test it?
|
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 |
|
A quick update, since part of my review above is out of date. JuliaDiff/ChainRulesCore.jl#724 no longer makes This PR is still useful. Vector cotangents still show up when a tuple is broadcast against a vector, e.g. On Since you're in favour of the What's left: the test comment "reduction not compiled" (it should either be removed or tested), and a rebase on |
# Conflicts: # Project.toml
- 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 <noreply@anthropic.com>
unbroadcast(x::Tuple, dx)only needs to reduce over trailing dimensions when a tuple was broadcast against something with more than one dimension:For a vector
dxthat branch is dead. A tuple of length N ≥ 2 (Tuple{Any}has its own method) fills the first dimension, solength(dx) == Nalways holds, andsum(dx; dims=2:1)would be a no-op anyway. Inference can't see this from the runtimeN == length(dx)check, though, so it still compilessum(dx; dims=...)for the eltype ofdx. In the generic broadcast path (split_bc_pullbacks) that eltype is often abstract.This PR checks
dx isa AbstractVectorfirst, so the branch is pruned statically. Runtime behaviour doesn't change.Minimal reproducer
First call, fresh session, Julia 1.13.3:
dxVector{Float64}Vector{Any}Vector{AbstractVector{Float64}}Downstream impact
The dead reduction becomes a real problem when inferring
+on the abstract eltype is expensive. In QuantumKitHub/BlockTensorKit.jl#86,Zygote.gradientofx -> (a, b = adjoint.(x); real(dot(a, a) + dot(b, b)))on a tuple of TensorKit tensors goes like this:Tangent{Any, Tuple{...}}cotangent, soback_genericcollects the per-element tangents into aVector{<:AdjointTensorMap{...}}with an abstract eltype.sumover 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 NTupletests;rulesets/Base/broadcastpasses locally. Version bumped to 1.74.9.Update (2026-11-12, by @devmotion)
This PR stopped receiving updates, so I merged
mainand finished it. The version is now 1.74.15 (#858 takes 1.74.14)._sum_trailing_dims, which dispatches on the cotangent type. Vectors,Tuples andTangents are returned as is, so inference never sees the reduction for them; other arrays are summed over2:ndims(dx).ndimsis only called onAbstractArrays.DimensionMismatch: a cotangent of the wrong length was silently truncated byTuple{Vararg{Any,N}}(val). It now throws.unbroadcast(x, @thunk(ZeroTangent()))failed for tuples (length(::ZeroTangent)) and arrays (ndims(::ZeroTangent)). Both now return early after unthunking. This makes the::AbstractZeromethods for tuples and arrays, and theAbstractArray{Bool}disambiguation method, redundant, so they are removed.Vector{Any}, N×1, matrix and 3-d arrays,Tuple,Tangent{Any}, thunked, zero and thunked-zero cotangents, and both length mismatches. They use@inferredwherever the result is inferable;Vector{Any}isn't, since its eltype could be a zero.First call in a fresh session (Julia 1.13.1,
mainvs this PR):dxmainVector{Float64}Vector{Any}Vector{AbstractVector{Float64}}🤖 Generated with Claude Code