Skip to content

Keep tuple structure of broadcast cotangents - #858

Closed
lkdvos wants to merge 1 commit into
JuliaDiff:mainfrom
lkdvos:tuple-tangent-broadcast
Closed

lkdvos wants to merge 1 commit into
JuliaDiff:mainfrom
lkdvos:tuple-tangent-broadcast

Conversation

@lkdvos

@lkdvos lkdvos commented Oct 29, 2026

Copy link
Copy Markdown

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

The cotangent of a tuple-valued broadcast arrives as a `Tangent`, which
Zygote constructs as `Tangent{Any}`. Broadcasting over it falls back to
`collect`, so the pullbacks produced `Vector`s with abstract element types
instead of tuples. Broadcast over the backing tuple instead.

Also make `unzip(::Tuple)` inferable by dispatching to the `Val{N}` method,
as is already done for arrays.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@codecov

codecov Bot commented Oct 29, 2026 •

Copy link
Copy Markdown

Codecov Report

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

Additional details and impacted files
@@           Coverage Diff           @@
##             main     #858   +/-   ##
=======================================
  Coverage   96.63%   96.63%           
=======================================
  Files          28       28           
  Lines        3596     3598    +2     
=======================================
+ Hits         3475     3477    +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

Copy link
Copy Markdown
Member

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

@devmotion

Copy link
Copy Markdown
Member

I opened JuliaDiff/ChainRulesCore.jl#724

@lkdvos

lkdvos commented Oct 29, 2026

Copy link
Copy Markdown
Author

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants