[lkdvos]

When a broadcast over tuples is differentiated, the cotangent of the (tuple) output arrives as a Tangent. Zygote builds it as Tangent{Any}, since ZBack doesn't see the primal. ChainRulesCore defines no broadcastable for Tangent, so the broadcast pullbacks hit Base's collect fallback. They then work with a Vector of whatever eltype collect finds, instead of a tuple:

julia> using ChainRules, ChainRulesCore

julia> Base.Broadcast.broadcastable(Tangent{Any}(1.0, [2.0]))
2-element Vector{Any}:
 1.0
  [2.0]

The element types are lost, and downstream code gets inferred on abstract eltypes. In addition, unzip(::Tuple) isn't inferable: ntuple(i -> map(Get(i), xs), length(x1)) constructs Get(i) from a runtime i.

This PR:

• broadcasts over the backing tuple of a Tangent{<:Any, <:Tuple} cotangent, in all four broadcast pullbacks;
• makes unzip(::Tuple) dispatch to the existing unzip(xs, ::Val{N}) generated method, as the array path already does.

Runtime results don't change; inference improves:

using ChainRules, ChainRulesCore
using Base.Broadcast: broadcasted
struct Cfg <: RuleConfig{Union{HasReverseMode}} end
ChainRulesCore.rrule_via_ad(::Cfg, f, args...) = rrule(f, args...)

_, bk = rrule(Cfg(), copy∘broadcasted, Broadcast.BroadcastStyle(Tuple), log, (1.0, 2.0f0))
Base.infer_return_type(bk, Tuple{Tangent{Any, Tuple{Float64, Float32}}})
# before: Tuple{NoTangent, NoTangent, Any, Vararg{Any}}
# after:  Tuple{NoTangent, NoTangent, NoTangent, Tangent{Tuple{Float64, Float32}, Tuple{Float64, Float32}}}

Base.infer_return_type(ChainRules.unzip, Tuple{Tuple{Tuple{Int, Float64}, Tuple{Int, Float64}}})
# before: Tuple{Tuple{Any, Any}, Tuple{Any, Any}}
# after:  Tuple{Tuple{Int64, Int64}, Tuple{Float64, Float64}}

Downstream impact

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 through the generic path (split_bc_pullbacks). The per-element tangents ended up in a Vector{<:AdjointTensorMap{...}} with an abstract eltype, and inferring the downstream code on that type took minutes once BlockTensorKit was loaded. With this PR, back_generic infers concrete tangent types and the first gradient takes 5.5 s (>150 s before; 6.7 s without BlockTensorKit loaded).

This is independent of #857, which avoids inferring the dead reduction in unbroadcast(::Tuple, ::AbstractVector). That case still occurs for a tuple broadcast against a vector, e.g. (a, b) .* [1, 2]. Both PRs bump the version to 1.74.9, so whichever is merged second needs a trivial rebase.

Tests

• Tangent{Any} cotangents for tuple broadcasts on split paths 2 and 4, including an @inferred check for path 2.
• @inferred tests for unzip / unzip_broadcast on tuples.

The @inferred tests fail on main. The full test suite passes locally on Julia 1.13.3.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.63%. Comparing base (a441bd5) to head (6b8be24).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #858   +/-   ##
=======================================
  Coverage   96.63%   96.63%           
=======================================
  Files          28       28           
  Lines        3598     3600    +2     
=======================================
+ Hits         3477     3479    +2     
  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.

[devmotion]

AFAICT ChainRules is the wrong package for this change, it's rather that ChainRulesCore is missing a broadcastable definition for Tangent.

[devmotion]

I opened JuliaDiff/ChainRulesCore.jl#724

[lkdvos]

Thanks for picking this up! I have to say I'm not too familiar with where these things go so I do appreciate the help

[devmotion]

I'm reopening this because I was wrong when I closed it: ChainRules is the right place for this change after all.

I tried to add a broadcastable definition for Tangent to ChainRulesCore in JuliaDiff/ChainRulesCore.jl#724, but no definition there is correct for all Tangents:

• With P === Any, as Zygote produces, the primal type is unknown, so ChainRulesCore cannot tell whether the tangent should act as a container or a scalar.
• A structural tangent of an array-like primal, such as Tangent{SVector}(data=...), would have to be a container, so treating it as a scalar is wrong. ChainRules' pullback then silently returns NoTangent()s.
• Even a tuple-backed tangent has two possible shapes. In a tuple-valued broadcast it is a container, but in a 0-dim broadcast it is a single element.

Only the code doing the broadcast knows which shape is meant, here the broadcast rule, which has the primal output. So #724 now makes broadcasting over structural tangents an ArgumentError, like Base does for NamedTuple and Dict, rather than silently collecting them. That means this PR is no longer just an inference improvement: the broadcast rules must unwrap the cotangent before #724 can be released.

What's left here: rebase on main, since the unzip change and the version bump conflict with the merged #859. Ideally, also run the tests against JuliaDiff/ChainRulesCore.jl#724.

[devmotion]

I reran this branch against JuliaDiff/ChainRulesCore.jl#724. rulesets/Base/broadcast passes, and Zygote gradients through log.(x) and (sin∘cos).(x) work for tuples. The four split paths look correct to me, and a thunked cotangent in back_generic is fine too, because back_generic(dys::AbstractThunk) unthunks first.

However, the fused rules broadcast over the cotangent as well, and they accept tuples: NumericOrBroadcast includes NTuple{<:Any,Number}. With #724 they throw for Tangent cotangents. With ChainRulesCore 1.26.8 they work:

julia> Zygote.gradient(t -> (y = t .* 2; y[1] + y[2]), (1.0, 2.0f0))
ERROR: ArgumentError: broadcasting over structural tangents is not supported; use `map` or arithmetic

The same happens for t .^ 2, t ./ 2 and imag.(t). .+, .-, float., T., conj. and complex. are fine, because they don't broadcast over the cotangent. I checked locally that unwrapping the cotangent in these four rules fixes it:

     function bc_times_back(Δraw)
-        Δ = unthunk(Δraw)
+        Δ = _bc_cotangent(unthunk(Δraw))
@@
     function bc_square_back(dy_raw)
-        dx = @thunk ProjectTo(x)(2 .* unthunk(dy_raw) .* conj.(x))
+        dx = @thunk ProjectTo(x)(2 .* _bc_cotangent(unthunk(dy_raw)) .* conj.(x))
@@
     function bc_divide_back(dz_raw)
-        dz = unthunk(dz_raw)
+        dz = _bc_cotangent(unthunk(dz_raw))
@@
-    bc_imag_back(dz) = (NoTangent(), NoTangent(), @thunk(im .* real.(unthunk(dz))))
+    bc_imag_back(dz) = (NoTangent(), NoTangent(), @thunk(im .* real.(_bc_cotangent(unthunk(dz)))))

The test suite passes on #724 only because the fused-rule tests don't cover tuples for these rules. Could you add tests like these to the "fused rules" testset? With #724, each of them errors without the change above and passes with it:

test_rrule(copy∘broadcasted, *, Tuple(rand(3)), rand())
test_rrule(copy∘broadcasted, *, rand(), Tuple(rand(3)))
test_rrule(copy∘broadcasted, *, Tuple(rand(3)), Tuple(rand(3)))
test_rrule(copy∘broadcasted, Base.literal_pow, ^, Tuple(rand(3)), Val(2))
test_rrule(copy∘broadcasted, /, Tuple(rand(3)), rand())

imag with a tuple fails test_rrule with or without the change, because bc_imag_back returns an unprojected Tuple (a Vector on main). That's a separate issue, so I left it out of the list above.

Also, main is already at 1.74.9 from #859, which hasn't been released yet, so VersionVigilante wants 1.74.10.

[devmotion]

With JuliaDiff/ChainRulesCore.jl#724, broadcasting over a Tangent throws instead of collecting it. ChainRules supports both versions, so maybe:

Suggested change

	# The cotangent of a tuple-valued broadcast is a `Tangent` (`Tangent{Any}` from Zygote). Broadcasting
	# over that would `collect` it into a `Vector`, losing the element types, so use its backing tuple.
	# The cotangent of a tuple-valued broadcast can be a `Tangent`

[devmotion]

Same as above:

Suggested change

	# Zygote passes the cotangent of a tuple as `Tangent{Any}`, which broadcasting used to
	# `collect` into a `Vector`, losing the element types.

[devmotion]

Maybe also cover thunked cotangents and the multi-argument path (bc_many_back), which is also inferable:

Suggested change

	@test all(ChainRulesCore.backing(dx4[4]) .≈ (-cos(cos(1.0)) * sin(1.0), -cos(cos(2.0f0)) * sin(2.0f0)))
	@test all(ChainRulesCore.backing(dx4[4]) .≈ (-cos(cos(1.0)) * sin(1.0), -cos(cos(2.0f0)) * sin(2.0f0)))
	@test bk2(@thunk(Tangent{Any}(1.0, 1.0f0))) == dx2
	@test bk4(@thunk(Tangent{Any}(1.0, 1.0f0))) == dx4
	_, bkm = rrule(CFG, copy∘broadcasted, BT1, atan, x, x) # split 2, several arguments
	dxm = @inferred bkm(Tangent{Any}(1.0, 1.0f0))
	@test dxm[4] isa Tangent{typeof(x), typeof(x)}
	@test dxm[5] isa Tangent{typeof(x), typeof(x)}