Repository navigation
Conversation
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 Report✅ All modified and coverable lines are covered by tests. 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. 🚀 New features to boost your workflow:
|
Member
|
AFAICT ChainRules is the wrong package for this change, it's rather that ChainRulesCore is missing a |
Member
|
I opened JuliaDiff/ChainRulesCore.jl#724 |
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 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When a broadcast over tuples is differentiated, the cotangent of the (tuple) output arrives as a
Tangent. Zygote builds it asTangent{Any}, sinceZBackdoesn't see the primal. ChainRulesCore defines nobroadcastableforTangent, so the broadcast pullbacks hit Base'scollectfallback. They then work with aVectorof whatever eltypecollectfinds, instead of a tuple: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))constructsGet(i)from a runtimei.This PR:
Tangent{<:Any, <:Tuple}cotangent, in all four broadcast pullbacks;unzip(::Tuple)dispatch to the existingunzip(xs, ::Val{N})generated method, as the array path already does.Runtime results don't change; inference improves:
Downstream impact
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 through the generic path (split_bc_pullbacks). The per-element tangents ended up in aVector{<:AdjointTensorMap{...}}with an abstract eltype, and inferring the downstream code on that type took minutes once BlockTensorKit was loaded. With this PR,back_genericinfers 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@inferredcheck for path 2.@inferredtests forunzip/unzip_broadcaston tuples.The
@inferredtests fail onmain. The full test suite passes locally on Julia 1.13.3.🤖 Generated with Claude Code