From 2652fa445ac69d3f28c87442cae51016ece8f03e Mon Sep 17 00:00:00 2001 From: lkdvos Date: Thu, 29 Oct 2026 14:59:30 -0400 Subject: [PATCH] Keep tuple structure of broadcast cotangents 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 --- Project.toml | 2 +- src/rulesets/Base/broadcast.jl | 13 +++++++++---- src/unzipped.jl | 2 +- test/rulesets/Base/broadcast.jl | 14 ++++++++++++++ test/unzipped.jl | 5 +++++ 5 files changed, 30 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index 2a59656eb..a18c5c4b0 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.8" +version = "1.74.9" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Base/broadcast.jl b/src/rulesets/Base/broadcast.jl index ae389b84f..726c1e031 100644 --- a/src/rulesets/Base/broadcast.jl +++ b/src/rulesets/Base/broadcast.jl @@ -2,6 +2,11 @@ using Base.Broadcast: Broadcast, broadcasted, Broadcasted, BroadcastStyle const RCR = RuleConfig{>:HasReverseMode} const TRI_NO = (NoTangent(), NoTangent(), NoTangent()) +# 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. +_bc_cotangent(dys::Tangent{<:Any, <:Tuple}) = ChainRulesCore.backing(dys) +_bc_cotangent(dys) = dys + function rrule(::typeof(copy), bc::Broadcasted) uncopy(Δ) = (NoTangent(), Δ) return copy(bc), uncopy @@ -59,7 +64,7 @@ function split_bc_derivatives(f::F, arg) where {F} @debug("split broadcasting derivative", f) ys = f.(arg) function bc_one_back(dys) # For f.(x) we do not need StructArrays / unzip at all - delta = broadcast(unthunk(dys), ys, arg) do dy, y, a + delta = broadcast(_bc_cotangent(unthunk(dys)), ys, arg) do dy, y, a das = only(derivatives_given_output(y, f, a)) dy * conj(only(das)) # possibly this * should be made nan-safe. end @@ -72,7 +77,7 @@ function split_bc_derivatives(f::F, args::Vararg{Any,N}) where {F,N} @debug("split broadcasting derivatives", f, N) ys = f.(args...) function bc_many_back(dys) - deltas = unzip_broadcast(unthunk(dys), ys, args...) do dy, y, as... + deltas = unzip_broadcast(_bc_cotangent(unthunk(dys)), ys, args...) do dy, y, as... das = only(derivatives_given_output(y, f, as...)) map(da -> dy * conj(da), das) # possibly this * should be made nan-safe. end @@ -110,7 +115,7 @@ function split_bc_inner(frule_fun::R, cfg::RuleConfig, f::F, arg) where {R,F} frule_fun(cfg, (NoTangent(), one(a)), f, a) end function back_forwards(dys) - delta = broadcast(ydots, unthunk(dys), arg) do ydot, dy, a + delta = broadcast(ydots, _bc_cotangent(unthunk(dys)), arg) do ydot, dy, a ProjectTo(a)(conj(ydot) * dy) # possibly this * should be made nan-safe. end return (TRI_NO..., ProjectTo(arg)(delta)) @@ -129,7 +134,7 @@ function split_bc_pullbacks(cfg::RCR, f::F, args::Vararg{Any,N}) where {F,N} rrule_via_ad(cfg, f, a...) end function back_generic(dys) - deltas = unzip_broadcast(backs, dys) do back, dy # (could be map, sizes match) + deltas = unzip_broadcast(backs, _bc_cotangent(dys)) do back, dy # (could be map, sizes match) map(unthunk, back(dy)) end dargs = map(unbroadcast, args, Base.tail(deltas)) diff --git a/src/unzipped.jl b/src/unzipped.jl index fe5875e6f..4bd9e6d63 100644 --- a/src/unzipped.jl +++ b/src/unzipped.jl @@ -132,7 +132,7 @@ julia> unzip(((1,2), (30,40), (500,600))) function unzip(xs::Tuple) x1 = first(xs) x1 isa Tuple || throw(ArgumentError("unzip only accepts arrays or tuples of tuples")) - return ntuple(i -> map(Get(i), xs),length(x1)) + return unzip(xs, Val(length(x1))) end struct Get{i} end diff --git a/test/rulesets/Base/broadcast.jl b/test/rulesets/Base/broadcast.jl index ea9bac8d0..ad37cd5c5 100644 --- a/test/rulesets/Base/broadcast.jl +++ b/test/rulesets/Base/broadcast.jl @@ -197,6 +197,20 @@ BT1 = Broadcast.BroadcastStyle(Tuple) @test y3 == () @test bk3(Tangent{Tuple{}}()) == (NoTangent(), NoTangent(), NoTangent(), ZeroTangent(), ZeroTangent()) end + @testset "tuple broadcast with Tangent{Any} cotangent" begin + # Zygote passes the cotangent of a tuple as `Tangent{Any}`, which broadcasting used to + # `collect` into a `Vector`, losing the element types. + x = (1.0, 2.0f0) + y2, bk2 = rrule(CFG, copy∘broadcasted, BT1, log, x) # split 2 + dx2 = @inferred bk2(Tangent{Any}(1.0, 1.0f0)) + @test dx2[4] isa Tangent{typeof(x), typeof(x)} + @test all(ChainRulesCore.backing(dx2[4]) .≈ (1.0, 0.5)) + + y4, bk4 = rrule(CFG, copy∘broadcasted, BT1, sin∘cos, x) # split 4 + dx4 = bk4(Tangent{Any}(1.0, 1.0f0)) + @test dx4[4] isa Tangent{typeof(x), typeof(x)} + @test all(ChainRulesCore.backing(dx4[4]) .≈ (-cos(cos(1.0)) * sin(1.0), -cos(cos(2.0f0)) * sin(2.0f0))) + end @testset "unbroadcast with NTuple" begin # https://github.com/JuliaDiff/ChainRules.jl/pull/661 @test ChainRules.unbroadcast((1, 2, [3]), [4, 5, [6]]) isa Tangent # earlier, NTuple demanded same type @test ChainRules.unbroadcast(broadcasted(-, (1, 2), 3), (4, 5)) == (4, 5) # earlier, called ndims(::Tuple) diff --git a/test/unzipped.jl b/test/unzipped.jl index 6bc508f3f..4bc8e4c77 100644 --- a/test/unzipped.jl +++ b/test/unzipped.jl @@ -27,6 +27,11 @@ using ChainRules: unzip_broadcast, unzip #, unzip_map @test fun(tuple, (1,2,3), [4,5,6]) == ([1, 2, 3], [4, 5, 6]) # mix tuple & vector end + @testset "unzip tuple of tuples is inferred" begin + @test @inferred(unzip(((1, 2.0), (3, 4.0)))) == ((1, 3), (2.0, 4.0)) + @test @inferred(unzip_broadcast(x -> (x, 2x), (1, 2.0))) == ((1, 2.0), (2, 4.0)) + end + @testset "rrules" begin # These exist to allow for second derivatives