From b234e7dd6ff92a2279726d929bd7ce6f70de4215 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 22 Sep 2026 14:12:30 +0200 Subject: [PATCH 01/36] Put broadcasting on Julia's documented custom-broadcast interface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Almost all of ReverseDiff's broadcasting hung off entry points Julia never documented as extension points, and the documented ones were unused: 2191 `Base.broadcast` overloads, 1777 `Base.broadcast!`, and a 2001-method `materialize` table keyed on *any* package's style, against five `copy` methods and no `copyto!`. Because `broadcast` was overloaded directly it short-circuited the pipeline before a `Broadcasted` was ever built, so `f.(x)` and `broadcast(f, x)` could diverge, which `Base` guarantees they cannot. `TrackedStyle` becomes `TrackedStyle{N} <: AbstractArrayStyle{N}`, so it tracks dimensionality and picks up `Base`'s scalar, tuple and `Unknown` precedence rules for free. The `Base.broadcast` overloads and the `materialize` table are deleted, making `f.(x)` and `broadcast(f, x)` the same computation by construction. `copyto!` now throws a descriptive error instead of failing after polluting the tape; `.=` with a tracked right-hand side has never worked. `copy` no longer asks `Core.Compiler.return_type` what a broadcast will produce in order to pick an implementation. The only question the argument types cannot answer is whether a tracked value reaches `f` other than through the arguments, such as a closure capturing one, and `mayhavetracked` already reports that. What the output turns out to be is settled after the forward pass instead: anything `f` hands back that is not a `Dual` does not depend on the tracked arguments, so it carries no partials and is left off the tape. That covers a `Bool` from a predicate, an `Int` from `exponent`, a non-`Real`, and a function constant in its argument, with one test. The `:tracker` implementation goes with it, along with three `@generated` functions. `∇broadcast` seeds `Dual`s only for the tracked arguments. `DiffRules` gives `besselj` a literal `NaN` partial for its order, and `NaN * 0.0` is `NaN`, so dualizing an untracked argument poisoned every partial of the same `Dual`. Narrowing the seeding also narrows the `Dual` to the number of tracked arguments and removes the `SVector`/`DiffResults` machinery. The tag is now derived once for the whole broadcast from the arguments that are actually seeded, rather than per element from the element types. Where `f`'s partials are known in closed form, `knownpartials` dispatches on the operation and its arguments and the cache keeps a `Contract(op, args)` per tracked argument instead of a `Dual` per element. The entry says that argument collects `op(seed, args...)` from every element of the output's derivative, with `args` naming broadcast arguments by position: `+`, `-`, `identity` `identity` or `-`, applied to the seed alone `*` `*` with the other argument `/`, `\` `/` by an untracked denominator Nothing is stored. A `Contract` is a zero-size type, and an argument a partial names is read back through `instruction.input`, where `record!` has already captured it, so the partial cannot go stale on a replay and an untracked array is not copied a second time. Every `knownpartials` method annotates all of `f`'s arguments, which is what makes naming one by position safe: those are exactly the arguments `splitargs` keeps, in order, so an `f` that also receives arguments held back from the tape matches no method. `op` meets the seed rather than the partial, so `/` forms no reciprocal, which is how `sum(y ./ 1e-310)` keeps the finite gradient that #298 gave it. A tracked denominator keeps a `Dual` per element, its partial `-x/y^2` being no argument of the broadcast, as does an argument whose shape is not the output's, which could not be read at the output's own index. Arrays of `TrackedReal` are differentiated like a `TrackedArray` rather than traced scalar-wise: 64x faster on `x .* y` and 166x on a fused three-argument expression, with the tape going from ~3n entries to 2. Against master, with `gradient` and Chairmarks, best of five alternating runs: case master.us branch.us t.ratio bytes.ratio mulscalar_1k 2.41 2.19 0.91 1.00 mulscalar_10k 15.75 14.33 0.91 1.00 mulscalar_100k 159.58 144.62 0.91 1.00 addscalar_10k 15.62 14.13 0.90 1.00 addarray_10k 17.33 15.75 0.91 1.00 mularray_10k 21.58 18.83 0.87 1.00 twoscalars_10k 24.67 23.00 0.93 0.87 sinscalar_10k 207.21 88.46 0.43 1.00 exp_10k 51.75 53.25 1.03 1.00 exp_100k 530.96 535.50 1.01 1.00 `exp` keeps a `Dual` per element either way and is at parity; the spread across runs is about 3%. Precompilation drops from 10.45s to 6.78s and load time from 2.61s to 2.39s, and `src` loses 417 lines. Fixes #265: a 0-d broadcast collapsed to a single `ImmutableDiffResult`, whose `length` was then taken. Fixes #214: a scalar and a one-element-vector argument now agree, where the scalar form raised a `ForwardOptimize` ambiguity. Addresses #226: the counted methods on `broadcast`, `broadcast!` and `materialize` go from 5969 to 0, and `Test.detect_ambiguities(ReverseDiff)` from 877 to 581, with the 296 broadcast-related ones all gone. `map`, `hcat` and `vcat` are untouched and still account for 4567. `seed!`'s catch-all is narrowed to `::Real`: it silently did nothing for the arrays and tuples that reach it from a broadcast whose output element type is not `Real`, and `jacobian` now covers those outputs. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 2 +- src/derivatives/broadcast.jl | 413 +++++++++++++---------- src/derivatives/elementwise.jl | 473 ++------------------------- src/derivatives/propagation.jl | 129 +++----- src/tracked.jl | 6 +- test/PrecisionTests.jl | 3 +- test/TrackedTests.jl | 13 +- test/api/JacobianTests.jl | 5 + test/derivatives/BroadcastTests.jl | 360 ++++++++++++++++++++ test/derivatives/ElementwiseTests.jl | 3 +- test/runtests.jl | 6 + 11 files changed, 684 insertions(+), 729 deletions(-) create mode 100644 test/derivatives/BroadcastTests.jl diff --git a/Project.toml b/Project.toml index c9057abd..058fe4cd 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ReverseDiff" uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267" -version = "1.18.2" +version = "1.19.0" [deps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index abecc67f..c456757f 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -2,10 +2,8 @@ ## Broadcasting ## ################## -using Base.Broadcast: BroadcastStyle, ArrayStyle, Broadcasted, broadcasted +using Base.Broadcast: AbstractArrayStyle, BroadcastStyle, Broadcasted, DefaultArrayStyle using ForwardDiff: ForwardDiff, Dual -import Base.Broadcast: materialize -const RDBroadcasted{F, T} = Broadcasted{<:Any, <:Any, F, T} """ NotTracked(f::Function) @@ -22,7 +20,6 @@ istypeorclosure(::AbstractArray{F}) where {F} = _istypeorclosure(F) istypeorclosure(::Base.RefValue{F}) where {F} = _istypeorclosure(F) istypeorclosure(::AbstractArray{<:Real}) = false istypeorclosure(::TrackedArray) = false -istypeorclosure(::AbstractArray{<:TrackedReal}) = true istypeorclosure(::Real) = false @generated function _istypeorclosure(::Type{F}) where {F} # `fieldcount` errors for types without a definite number of fields, such as @@ -38,25 +35,40 @@ end mayhavetracked(b) = istypeorclosure(b) mayhavetracked(b::Type) = false mayhavetracked(b::NotTracked) = false +mayhavetracked(b::ForwardOptimize) = mayhavetracked(b.f) +mayhavetracked(b::SkipOptimize) = mayhavetracked(b.f) mayhavetracked(b::Base.RefValue{<:NotTracked}) = false mayhavetracked(b::Broadcasted) = mayhavetracked(b.f) || any(mayhavetracked, b.args) -struct TrackedStyle <: BroadcastStyle end +struct TrackedStyle{N} <: AbstractArrayStyle{N} end -Broadcast.BroadcastStyle(::Type{<:Union{TrackedArray, TrackedReal}}) = TrackedStyle() -Broadcast.BroadcastStyle(::TrackedStyle, b::BroadcastStyle) = TrackedStyle() +(::Type{<:TrackedStyle})(::Val{N}) where {N} = TrackedStyle{N}() -# We have to re-build the original broadcast struct to get the appropriate array -# style. We need this primarily to support CuArrays' broadcasting fixes. -broadcast_rebuild(xs) = recur_value(xs) -recur_value(xs) = xs -recur_value(xs::Union{TrackedReal, TrackedArray}) = recur_value(value(xs)) +Broadcast.BroadcastStyle(::Type{<:TrackedArray{V,D,N}}) where {V,D,N} = TrackedStyle{N}() +Broadcast.BroadcastStyle(::Type{<:TrackedReal}) = TrackedStyle{0}() +Broadcast.BroadcastStyle(::Type{<:AbstractArray{<:TrackedReal,N}}) where {N} = TrackedStyle{N}() -function broadcast_rebuild(bc::Broadcasted) - broadcasted(bc.f, broadcast_rebuild.(bc.args)...) -end +# `AbstractArrayStyle{Any}` carries `Any` as its dimension, which `max` cannot compare +_maxdim(M::Int, N::Int) = max(M, N) +_maxdim(::Type{Any}, ::Int) = Any +_maxdim(::Int, ::Type{Any}) = Any +_maxdim(::Type{Any}, ::Type{Any}) = Any + +# tracked values must stay tracked, so take precedence over every other array style +Broadcast.BroadcastStyle(::TrackedStyle{M}, ::AbstractArrayStyle{N}) where {M,N} = + TrackedStyle{_maxdim(M, N)}() + +# resolve the overlap with `Base`'s three `DefaultArrayStyle` rules, preserving their results +Broadcast.BroadcastStyle(::TrackedStyle{M}, ::DefaultArrayStyle{N}) where {M,N} = + TrackedStyle{_maxdim(M, N)}() +Broadcast.BroadcastStyle(::TrackedStyle{N}, ::DefaultArrayStyle{N}) where {N} = TrackedStyle{N}() +Broadcast.BroadcastStyle(::TrackedStyle{Any}, ::DefaultArrayStyle) = TrackedStyle{Any}() + +# the untracked style decides where the fallback re-dispatches, and a `CuArray` backing a +# `TrackedArray` has to keep its own +recur_value(xs) = xs +recur_value(xs::Union{TrackedReal, TrackedArray, AbstractArray{<:TrackedReal}}) = recur_value(value(xs)) -getstyle(::Broadcasted{Style}) where {Style} = Style remove_not_tracked(f) = f remove_not_tracked(f::NotTracked) = f.f remove_not_tracked(f::Base.RefValue{<:NotTracked}) = Ref(remove_not_tracked(f[])) @@ -65,59 +77,39 @@ function remove_not_tracked(b::Broadcasted{style}) where {style} return Broadcasted{style}(remove_not_tracked(b.f), remove_not_tracked.(b.args), b.axes) end -onlyrealarrays(args::Tuple) = onlyrealarray(first(args)) && onlyrealarrays(Base.tail(args)) -onlyrealarrays(::Tuple{}) = true -onlyrealarray(::AbstractArray{<:Real}) = true -onlyrealarray(::AbstractArray) = false -onlyrealarray(::Any) = true - -anyreals(args::Tuple) = first(args) isa Real || anyreals(Base.tail(args)) -anyreals(args::Tuple{}) = false - -function get_implementation(bc, f, T, args) - outputisreal = (T <: AbstractArray{<:Real}) && (T !== Union{}) - # Each arg is either a real number, an array of untraked reals, a tracked array of reals or an array of untracked non-reals, - # Output is real, and - # No tracked closure or arguments, except TrackedReal and TrackedArray. - if !mayhavetracked(bc) && outputisreal && (anyreals(args) || !onlyrealarrays(args)) - return Val(:tracker) - # No arg is a real number and array args must be arrays of untracked reals or tracked arrays of reals, - # Output is real, and - # No tracked closure or arguments, except TrackedReal and TrackedArray. - elseif !mayhavetracked(bc) && outputisreal - return Val(:reversediff) - # Function or any arg is possibly a tracked non-real or an array of tracked reals/non-reals, - # Or output is not an array of reals - else - return Val(:fallback) - end -end -function Base.copy(_bc::Broadcasted{TrackedStyle}) +# scalars take `Base`'s 0-dimensional route onto the scalar derivative rules; `instantiate` +# leaves an `AbstractArrayStyle{0}` without axes, so 0 dimensions shows up either way +Base.copy(bc::Broadcasted{<:TrackedStyle, <:Union{Nothing, Tuple{}}}) = bc[CartesianIndex()] + +function Base.copy(_bc::Broadcasted{<:TrackedStyle}) bc = remove_not_tracked(_bc) flattened_bc = Base.Broadcast.flatten(bc) - untracked_bc = broadcast_rebuild(bc) - T = Core.Compiler.return_type(copy, Tuple{typeof(untracked_bc)}) f, args = flattened_bc.f, flattened_bc.args - implementation = get_implementation(_bc, f, T, args) - if implementation isa Val{:reversediff} - return ∇broadcast(f, args...) - elseif implementation isa Val{:tracker} - return tracker_∇broadcast(f, args...) + # only the arguments are seeded, so a tracked value reaching `f` by another route, such + # as a closure capturing one, has to be traced scalar-wise + if mayhavetracked(_bc) + axs = flattened_bc.axes + style = typeof(Broadcast.combine_styles(map(recur_value, args)...)) + return copy(Broadcasted{style, typeof(axs), typeof(f), typeof(args)}(f, args, axs)) else - flattened_untracked_bc = Base.Broadcast.flatten(untracked_bc) - style, axes = getstyle(flattened_untracked_bc), flattened_bc.axes - return copy(Broadcasted{style, typeof(axes), typeof(f), typeof(args)}(f, args, axes)) + return ∇broadcast(f, args...) end end +_no_tracked_dest() = throw(ArgumentError("`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead.")) + +Base.copyto!(::TrackedArray, ::Broadcasted{<:TrackedStyle}) = _no_tracked_dest() +Base.copyto!(::TrackedArray, ::Broadcasted{<:DefaultArrayStyle}) = _no_tracked_dest() + getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::TrackedArray{<:Any, D}) where {D} = D +getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D getouttype(::Any) = Union{} deref(x) = x deref(x::Base.RefValue) = x[] -@generated function splatcall(f, x::SVector{N}, utargs::T, ::Val{tinds}) where {N, T <: Tuple, tinds} +@generated function splatcall(f, x::NTuple{N,Any}, utargs::T, ::Val{tinds}) where {N, T <: Tuple, tinds} args = [] ti = 1 uti = 1 @@ -148,157 +140,220 @@ end ## A generalization of the broadcasting approach in ReverseDiff for general functions -@inline function ∇broadcast(f::F, args::Vararg{Any}) where {F} - inds, targs, untracked = splitargs(args) - N = length(targs) - D = promote_type(getouttype.(targs)...) - result = DiffResults.GradientResult(zero(SVector{N, D})) - function df(x...) - return ForwardDiff.gradient!( - result, - s -> splatcall(f, s, untracked, inds), - SVector(x), - ) - end - results = broadcast(df, value.(targs)...) - tp = tape(targs...) - out_value = DiffResults.value.(results) - eltype(out_value) == Bool && return out_value - out = track(out_value, D, tp) - cache = (results, df, index_bound.(targs, (out,))) - record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache) - return out -end -@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) - input = instruction.input - output = instruction.output - output_deriv = deriv(output) - results, _, bounds = instruction.cache - N = length(input) - if N == 1 || all(isequal(size(input[1])), size.(Base.tail(input))) - _br_add_to_deriv!(input, output_deriv, results) - else - _br_add_to_deriv!(input, output_deriv, results, bounds) - end - unseed!(output) - return nothing +@inline incr(::Val{k}) where {k} = Val(k + 1) + +# `slots[i]` indexes argument `i`'s partial and is `Val(0)` where the argument is untracked; +# the second value counts the tracked arguments. Both are built on the way out of the +# recursion: a count passed into it would stop being constant-folded after a few arguments, +# leaving the partials a runtime-sized tuple. +@inline trackedslots(::Tuple{}) = (), Val(0) +@inline function trackedslots(args::Tuple) + slots, n = trackedslots(Base.tail(args)) + istracked(first(args)) || return (Val(0), slots...), n + k = incr(n) + return (k, slots...), k end -@generated function _br_add_to_deriv!(xs::T, o, r) where {T <: Tuple} - N = length(T.types) - return Expr(:block, [:(_br_add_to_deriv!(xs[$i], o, r, Val($i))) for i in 1:N]...) -end -_br_add_to_deriv!(_, _, _, _) = nothing -function _br_add_to_deriv!(x::Union{TrackedReal, TrackedArray}, out_deriv, results, ::Val{i}) where {i} - return istracked(x) && diffresult_increment_deriv!(x, out_deriv, results, i) +# `DiffRules` leaves partials such as the order of `besselj` as `NaN`, which `NaN * 0` +# would spread to every slot, so untracked arguments are not dualized +@inline dualize(::Type, ::Val{0}, ::Val, x) = x +@inline function dualize(::Type{T}, ::Val{k}, valP::Val, x) where {T, k} + return Dual{T}(x, ntuple(j -> j == k, valP)) end -@generated function _br_add_to_deriv!(xs::T, o, r, bounds) where {T <: Tuple} - N = length(T.types) - return Expr(:block, [:(_br_add_to_deriv!(xs[$i], o, r, Val($i), bounds[$i])) for i in 1:N]...) -end -_br_add_to_deriv!(_, _, _, _, _) = nothing -function _br_add_to_deriv!(x::Union{TrackedReal,TrackedArray}, out_deriv, results, ::Val{i}, bound) where {i} - return istracked(x) && diffresult_increment_deriv!(x, out_deriv, results, i, bound) +# only seeded arguments make up the `Dual` value type, so the tag follows `slots` +@inline dualvaltype(::Val{0}, v) = Union{} +@inline dualvaltype(::Val, v) = eltype(v) + +# `f`'s partials in closed form, one entry per tracked argument in slot order +struct KnownPartials{E<:Tuple} + entries::E end -@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) - input, output = instruction.input, instruction.output - results, df, _ = instruction.cache - pull_value!.(input) - broadcast!(df, results, value.(input)...) - output_value = value(output) - output_value .= DiffResults.value.(results) - return nothing +# `trackedslots` numbers the tracked arguments from the last, so drop the untracked ones and +# reverse what remains +trackedentries(::Tuple{}, ::Tuple{}) = () +trackedentries(slots::Tuple{Val{0},Vararg{Any}}, entries::Tuple) = + trackedentries(Base.tail(slots), Base.tail(entries)) +trackedentries(slots::Tuple, entries::Tuple) = + (trackedentries(Base.tail(slots), Base.tail(entries))..., first(entries)) + +# an argument read at the output's own index has to span the whole output +_sameshape(x, y) = x isa Real || y isa Real || axes(x) == axes(y) + +# one entry per argument, or `nothing` where `f`'s partials depend on the point. Annotating +# every argument is what lets an entry name one: those are exactly the ones `splitargs` keeps. +knownpartials(f, args...) = nothing + +# `+` and `-` are affine in their arguments jointly +knownpartials(::Union{typeof(+), typeof(identity)}, + args::Vararg{Union{Real, AbstractArray{<:Real}}}) = + map(_ -> Contract(identity, ()), args) + +knownpartials(::typeof(-), arg::Union{Real, AbstractArray{<:Real}}) = (Contract(-, ()),) + +knownpartials(::typeof(-), x::Union{Real, AbstractArray{<:Real}}, + y::Union{Real, AbstractArray{<:Real}}) = + (Contract(identity, ()), Contract(-, ())) + +# `*`, `/` and `\` are multilinear, so a partial is another argument, which the reverse pass +# reads from the instruction's own input where `record!` has captured it. A tracked +# denominator is excluded, its partial `-x/y^2` being no argument of the broadcast. +function knownpartials(::typeof(*), x::Union{Real, AbstractArray{<:Real}}, + y::Union{Real, AbstractArray{<:Real}}) + if _sameshape(x, y) + return (Contract(*, (Val(2),)), Contract(*, (Val(1),))) + else + return nothing + end end -## Tracker style broadcasting -## Good for broadcasting real numbers or arrays of non-tracked structs +function knownpartials(::typeof(/), x::Union{Real, AbstractArray{<:Real}}, + y::Union{Real, AbstractArray{<:Real}}) + if !istracked(y) && _sameshape(x, y) + return (Contract(/, (Val(2),)), nothing) + else + return nothing + end +end -trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x)))) +function knownpartials(::typeof(\), x::Union{Real, AbstractArray{<:Real}}, + y::Union{Real, AbstractArray{<:Real}}) + if !istracked(x) && _sameshape(x, y) + return (nothing, Contract(/, (Val(1),))) + else + return nothing + end +end -unbroadcast(x::AbstractArray, Δ) = - size(x) == size(Δ) ? Δ : - length(x) == length(Δ) ? trim(x, Δ) : - trim(x, sum(Δ, dims = ntuple(i -> size(x, i) == 1 ? i : ndims(Δ)+1, Val(ndims(Δ))))) +broadcastresults(::Nothing, slots, df, vals) = broadcast(df, vals...) +broadcastresults(entries::Tuple, slots, df, vals) = + KnownPartials(trackedentries(slots, entries)) -unbroadcast(x::Number, Δ) = sum(Δ) -unbroadcast(x::Base.RefValue, _) = nothing +# at least one argument has to be a non-0-dimensional array: `copy` sends the scalar and +# 0-dimensional cases down `Base`'s `TrackedStyle{0}` route onto the scalar rules instead +@inline function ∇broadcast(f::F, args::Vararg{Any}) where {F} + inds, targs, untracked = splitargs(args) + D = mapreduce(getouttype, promote_type, targs) + slots, valP = trackedslots(targs) + vals = map(value, targs) + # one tag for the whole broadcast keeps `results` concretely typed + T = typeof(ForwardDiff.Tag(f, reduce(promote_type, map(dualvaltype, slots, vals)))) + # `broadcast` calls `df` elementwise, so it receives one scalar per argument + function df(x::Vararg{Any,N}) where {N} + dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) + return splatcall(f, dx, untracked, inds) + end + # known partials leave nothing to read off a `Dual`, so `f` is evaluated undualized + vf(x::Vararg{Any,N}) where {N} = splatcall(f, x, untracked, inds) + entries = knownpartials(f, args...) + return trackresults(T, broadcastresults(entries, slots, df, vals), df, vf, targs, D) +end -dual(x, p) = x -dual(x::Real, p) = Dual(x, p) +# the cache carries what the replay needs: `df` to recompute the stored partials, or, where +# they are known already, `vf` for the values alone +replaycache(::Type{T}, results::AbstractArray, df, vf, targs) where {T} = + (df, map(y -> ForwardDiff.value(T, y), results)) +replaycache(::Type, ::KnownPartials, df, vf, targs) = + (vf, broadcast(vf, map(value, targs)...)) -function _deriv(f, G, ::Val{i}, args::Vararg{Any, N}) where {N, i} - dargs = ntuple(j -> dual(args[j], i==j), Val(N)) - return f(dargs...).partials[1] * G -end -@generated function _derivs(f, G, args::Vararg{Any, N}) where {N} - return Expr(:tuple, [:(_deriv.(f, G, Val($i), args...)) for i in 1:N]...) -end -@inline function tracker_∇broadcast(f, args::Vararg{Any, N}) where {N} - args_values = map(value, args) - out_value = broadcast(f, args_values...) - tp = tape(args...) - eltype(out_value) == Bool && return out_value - out = track(out_value, tp) - cache = (f,) - record!(tp, SpecialInstruction, tracker_∇broadcast, args, out, cache) +# `df` hands back `f`'s own result, so `results` keeps the element type `Base` would produce +@inline function recordresults(::Type{T}, results, df, vf, targs, ::Type{D}) where {T, D} + g, outvalue = replaycache(T, results, df, vf, targs) + tp = tape(targs...) + out = track(outvalue, D, tp) + cache = (results, g, T(), map(t -> index_bound(t, out), targs)) + record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache) return out end -@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)}) - input, output = instruction.input, instruction.output - f = instruction.cache[1] - output_value = value(output) - pull_value!.(input) - broadcast!(f, output_value, value.(input)...) +@inline trackresults(::Type{T}, results::AbstractArray{<:Dual{T}}, df, vf, targs, + ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, D) + +@inline trackresults(::Type{T}, results::KnownPartials, df, vf, targs, + ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, D) + +# an enclosing differentiation's tag is constant in our arguments, while one nested inside +# `f` buries our partial; a widened element type is a `Union`, so every member is checked +@inline function checktags(::Type{T}, ::Type{E}) where {T, E} + if E isa Union + checktags(T, E.a) + checktags(T, E.b) + else + S = ForwardDiff.tagtype(E) + if S !== Nothing && !ForwardDiff.:≺(S, T) + throw(ForwardDiff.DualMismatchError(T, S)) + end + end return nothing end -@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)}) +# a type-unstable `f`, such as one with an integer literal branch, leaves an abstract +# element type that can still hide a `Dual{T}` +@inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, + ::Type{D}) where {T, D} + if typeintersect(eltype(results), Dual{T}) === Union{} + checktags(T, eltype(results)) + return results + else + return recordresults(T, results, df, vf, targs, D) + end +end + +@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) input = instruction.input output = instruction.output - f = instruction.cache[1] output_deriv = deriv(output) - N = length(input) - Δargs = _derivs(f, output_deriv, value.(input)...) - dxs = map(unbroadcast, input, Δargs) - map(_add_to_deriv!, input, dxs) + results, _, tag, bounds = instruction.cache + T = typeof(tag) + slots, _ = trackedslots(input) + map((x, slot, bound) -> + _br_add_to_deriv!(T, x, slot, output_deriv, results, bound, input), + input, slots, bounds) unseed!(output) return nothing end -## Limited ReverseDiff broadcasting -## Efficient broadcasting for specific functions, e.g. +, - +# a per-element cache is indexed by the slot, a closed-form one is picked out by it +selectpartial(results, ::Val{k}, args) where {k} = (results, k) +selectpartial(p::KnownPartials, ::Val{k}, args) where {k} = (p.entries[k], args) -@inline _materialize(f, args) = broadcast(f, args...) +_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::CartesianIndex, _) = nothing +_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::Nothing, _) = nothing -for (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) - if !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f)) - @warn "$M.$f is not available and hence rule for it can not be defined" - continue # Skip rules for methods not defined in the current scope - end - if arity == 1 - @eval @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray}}) = _materialize(bc.f, bc.args) - elseif arity == 2 - @eval begin - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,TrackedArray}}) = _materialize(bc.f, bc.args) - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,TrackedReal}}) = _materialize(bc.f, bc.args) - @noinline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedReal,TrackedArray}}) = _materialize(bc.f, bc.args) - end - for A in ARRAY_TYPES - @eval begin - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A{<:Number},TrackedArray}}) = _materialize(bc.f, bc.args) - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray, $A{<:Number}}}) = _materialize(bc.f, bc.args) - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$A{<:Number}, TrackedReal}}) = _materialize(bc.f, bc.args) - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedReal,$A{<:Number}}}) = _materialize(bc.f, bc.args) - end - end - for R in REAL_TYPES - @eval begin - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{$R,TrackedArray}}) = _materialize(bc.f, bc.args) - @inline materialize(bc::RDBroadcasted{typeof($M.$f), <:Tuple{TrackedArray,$R}}) = _materialize(bc.f, bc.args) - end - end +# an argument broadcast to the full output shape needs no index clamping +function _br_add_to_deriv!(::Type{T}, x, slot::Val{k}, out_deriv, results, + bound::CartesianIndex, args) where {T, k} + results, sel = selectpartial(results, slot, args) + if bound == CartesianIndex(size(out_deriv)) + return diffresult_increment_deriv!(T, x, out_deriv, results, sel) + else + return diffresult_increment_deriv!(T, x, out_deriv, results, sel, bound) end end + +function _br_add_to_deriv!(::Type{T}, x, slot::Val{k}, out_deriv, results, ::Nothing, + args) where {T, k} + results, sel = selectpartial(results, slot, args) + return diffresult_increment_deriv!(T, x, out_deriv, results, sel, nothing) +end + +@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) + input, output = instruction.input, instruction.output + results, df, tag, _ = instruction.cache + foreach(pull_value!, input) + _replay!(typeof(tag), value(output), results, df, map(value, input)) + return nothing +end + +function _replay!(::Type{T}, out_value, results::AbstractArray, df, vals) where {T} + broadcast!(df, results, vals...) + map!(y -> ForwardDiff.value(T, y), out_value, results) + return nothing +end + +# known partials stay valid, so only the values are recomputed +function _replay!(::Type, out_value, results::KnownPartials, vf, vals) + broadcast!(vf, out_value, vals...) + return nothing +end diff --git a/src/derivatives/elementwise.jl b/src/derivatives/elementwise.jl index 130c7aa3..b56122a3 100644 --- a/src/derivatives/elementwise.jl +++ b/src/derivatives/elementwise.jl @@ -1,15 +1,15 @@ -############################### -# SkipOptimized map/broadcast # -############################### +##################### +# SkipOptimized map # +##################### # dispatch # #----------# -for g in (:map, :broadcast), f in SKIPPED_UNARY_SCALAR_FUNCS +for g in (:map,), f in SKIPPED_UNARY_SCALAR_FUNCS @eval @inline Base.$(g)(f::typeof($f), t::TrackedArray) = $(g)(SkipOptimize(f), t) end -for g in (:map, :broadcast), f in SKIPPED_BINARY_SCALAR_FUNCS +for g in (:map,), f in SKIPPED_BINARY_SCALAR_FUNCS @eval begin @inline Base.$(g)(f::typeof($f), x::TrackedArray, y::TrackedArray) = $(g)(SkipOptimize(f), x, y) @inline Base.$(g)(f::typeof($f), x::TrackedArray, y::TrackedReal) = $(g)(SkipOptimize(f), x, y) @@ -34,11 +34,11 @@ end # record # #--------# -for g in (:map, :broadcast) +for g in (:map,) @eval @inline Base.$(g)(f::SkipOptimize{F}, t::TrackedArray) where {F} = $(g)(f.f, value(t)) end -for g in (:map, :broadcast) +for g in (:map,) @eval begin @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedArray, y::TrackedArray) where {F} = $(g)(f.f, value(x), value(y)) @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedArray, y::TrackedReal) where {F} = $(g)(f.f, value(x), value(y)) @@ -60,136 +60,14 @@ for g in (:map, :broadcast) end end -#################################### -# ForwardOptimized map!/broadcast! # -#################################### +######################## +# ForwardOptimized map # +######################## # dispatch # #----------# -for g! in (:map!, :broadcast!), (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) - if !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f)) - @warn "$M.$f is not available and hence rule for it can not be defined" - continue # Skip rules for methods not defined in the current scope - end - (M, f) in SKIPPED_DIFFRULES && continue - if arity == 1 - @eval @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, t::TrackedArray) = $(g!)(ForwardOptimize(f), out, t) - elseif arity == 2 - @eval begin - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::TrackedArray, y::TrackedArray) = $(g!)(ForwardOptimize(f), out, x, y) - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::TrackedArray, y::TrackedReal) = $(g!)(ForwardOptimize(f), out, x, y) - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::TrackedReal, y::TrackedArray) = $(g!)(ForwardOptimize(f), out, x, y) - end - for A in ARRAY_TYPES - @eval begin - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::$A, y::TrackedArray) = $(g!)(ForwardOptimize(f), out, x, y) - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::TrackedArray, y::$A) = $(g!)(ForwardOptimize(f), out, x, y) - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::$A, y::TrackedReal) = $(g!)(ForwardOptimize(f), out, x, y) - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::TrackedReal, y::$A) = $(g!)(ForwardOptimize(f), out, x, y) - end - end - for R in REAL_TYPES - @eval begin - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::$R, y::TrackedArray) = $(g!)(ForwardOptimize(f), out, x, y) - @inline Base.$(g!)(f::typeof($M.$f), out::TrackedArray, x::TrackedArray, y::$R) = $(g!)(ForwardOptimize(f), out, x, y) - end - end - end -end - -# record # -#--------# - -for (g!, g) in ((:map!, :map), (:broadcast!, :broadcast)) - @eval function Base.$(g!)(f::ForwardOptimize{F}, out::TrackedArray{S}, x::TrackedArray{X}) where {F,S,X} - result = DiffResults.DiffResult(zero(S), zero(S)) - df = v -> ForwardDiff.derivative!(result, f.f, v) - results = $(g)(df, value(x)) - map!(DiffResult.value, value(out), results) - cache = (results, df, index_bound(x, out), nothing) - record!(tape(x), SpecialInstruction, $(g), x, out, cache) - return out - end - for TX in (:TrackedArray, :TrackedReal), TY in (:TrackedArray, :TrackedReal) - (TX == TY == :TrackedReal) && continue - @eval function Base.$(g!)(f::ForwardOptimize{F}, out::TrackedArray{S}, x::$(TX){X}, y::$(TY){Y}) where {F,S,X,Y} - result = DiffResults.GradientResult(SVector(zero(S), zero(S))) - df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) - results = $(g)(df, value(x), value(y)) - map!(DiffResult.value, value(out), results) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tape(x, y), SpecialInstruction, $(g), (x, y), out, cache) - return out - end - end - for A in ARRAY_TYPES - @eval function Base.$(g!)(f::ForwardOptimize{F}, out::TrackedArray, x::TrackedReal{X,D}, y::$A) where {F,X,D} - result = DiffResults.DiffResult(zero(X), zero(D)) - df = let result=result - (vx, vy) -> let vy=vy - ForwardDiff.derivative!(result, s -> f.f(s, vy), vx) - end - end - results = $(g)(df, value(x), value(y)) - map!(DiffResult.value, value(out), results) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tape(x), SpecialInstruction, $(g), (x, y), out, cache) - return out - end - @eval function Base.$(g!)(f::ForwardOptimize{F}, out::TrackedArray, x::$A, y::TrackedReal{Y,D}) where {F,Y,D} - result = DiffResults.DiffResult(zero(Y), zero(D)) - df = let result=result - (vx, vy) -> let vx=vx - ForwardDiff.derivative!(result, s -> f.f(vx, s), vy) - end - end - results = $(g)(df, value(x), value(y)) - map!(DiffResult.value, value(out), results) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tape(y), SpecialInstruction, $(g), (x, y), out, cache) - return out - end - @eval function Base.$(g!)(f::ForwardOptimize{F}, out::TrackedArray, x::TrackedArray{X}, y::$A) where {F,X} - result = DiffResults.GradientResult(SVector(zero(X))) - df = (vx, vy) -> let vy=vy - ForwardDiff.gradient!(result, s -> f.f(s[1], vy), SVector(vx)) - end - results = $(g)(df, value(x), value(y)) - map!(DiffResult.value, value(out), results) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tape(x), SpecialInstruction, $(g), (x, y), out, cache) - return out - end - @eval function Base.$(g!)(f::ForwardOptimize{F}, out::TrackedArray, x::$A, y::TrackedArray{Y}) where {F,Y} - result = DiffResults.GradientResult(SVector(zero(Y))) - df = let vx=vx - (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(vx, s[1]), SVector(vy)) - end - results = $(g)(df, value(x), value(y)) - map!(DiffResult.value, value(out), results) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tape(y), SpecialInstruction, $(g), (x, y), out, cache) - return out - end - end -end - -for R in REAL_TYPES - @eval begin - @inline Base.broadcast!(f::ForwardOptimize{F}, out::TrackedArray, x::TrackedArray, y::$R) where {F} = broadcast!(ForwardOptimize(t -> f.f(t, y)), out, x) - @inline Base.broadcast!(f::ForwardOptimize{F}, out::TrackedArray, x::$R, y::TrackedArray) where {F} = broadcast!(ForwardOptimize(t -> f.f(x, t)), out, y) - end -end - -################################## -# ForwardOptimized map/broadcast # -################################## - -# dispatch # -#----------# - -for g in (:map, :broadcast), (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) +for g in (:map,), (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) if !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f)) @warn "$M.$f is not available and hence rule for it can not be defined" continue # Skip rules for methods not defined in the current scope @@ -198,10 +76,6 @@ for g in (:map, :broadcast), (M, f, arity) in DiffRules.diffrules(; filter_modul @eval @inline Base.$(g)(f::typeof($M.$f), t::TrackedArray) = $(g)(ForwardOptimize(f), t) elseif arity == 2 (M, f) in SKIPPED_DIFFRULES && continue - # skip these definitions if `f` is one of the functions - # that will get a manually defined broadcast definition - # later (see "built-in infix operations" below) - g == :broadcast && in(f, (:+, :-, :*, :/, :\, :^)) && continue @eval begin @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::TrackedReal) = $(g)(ForwardOptimize(f), x, y) @@ -227,14 +101,14 @@ end # record # #--------# -for g in (:map, :broadcast) +for g in (:map,) @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedArray{X,D}) where {F,X,D} T = promote_type(X, D) result = DiffResults.DiffResult(zero(T), zero(T)) df = v -> ForwardDiff.derivative!(result, f.f, v) results = $(g)(df, value(x)) tp = tape(x) - out = track(DiffResults.value.(results), D, tp) + out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), nothing) record!(tp, SpecialInstruction, $(g), x, out, cache) return out @@ -249,7 +123,7 @@ for g in (:map, :broadcast) end results = $(g)(df, value(x), value(y)) tp = tape(x) - out = track(DiffResults.value.(results), D, tp) + out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), index_bound(y, out)) record!(tp, SpecialInstruction, $(g), (x, y), out, cache) return out @@ -263,7 +137,7 @@ for g in (:map, :broadcast) end results = $(g)(df, value(x), value(y)) tp = tape(y) - out = track(DiffResults.value.(results), D, tp) + out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), index_bound(y, out)) record!(tp, SpecialInstruction, $(g), (x, y), out, cache) return out @@ -275,7 +149,7 @@ for g in (:map, :broadcast) end results = $(g)(df, value(x), value(y)) tp = tape(x) - out = track(DiffResults.value.(results), D, tp) + out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), index_bound(y, out)) record!(tp, SpecialInstruction, $(g), (x, y), out, cache) return out @@ -287,7 +161,7 @@ for g in (:map, :broadcast) end results = $(g)(df, value(x), value(y)) tp = tape(y) - out = track(DiffResults.value.(results), D, tp) + out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), index_bound(y, out)) record!(tp, SpecialInstruction, $(g), (x, y), out, cache) return out @@ -301,7 +175,7 @@ for g in (:map, :broadcast) df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) results = $(g)(df, value(x), value(y)) tp = tape(x, y) - out = track(DiffResults.value.(results), D, tp) + out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), index_bound(y, out)) record!(tp, SpecialInstruction, $(g), (x, y), out, cache) return out @@ -309,18 +183,11 @@ for g in (:map, :broadcast) end end -for R in REAL_TYPES - @eval begin - @inline Base.broadcast(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::$R) where {F,X,D} = broadcast(ForwardOptimize(t -> f.f(t, y)), x) - @inline Base.broadcast(f::ForwardOptimize{F}, x::$R, y::TrackedArray{Y,D}) where {F,Y,D} = broadcast(ForwardOptimize(t -> f.f(x, t)), y) - end -end - ################ # forward pass # ################ -for (g!, g) in ((:map!, :map), (:broadcast!, :broadcast)) +for (g!, g) in ((:map!, :map),) @eval begin @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof($g)}) input, output = instruction.input, instruction.output @@ -352,315 +219,19 @@ end output_deriv = deriv(output) results = first(instruction.cache) if istracked(input) - diffresult_increment_deriv!(input, output_deriv, results, 1) + diffresult_increment_deriv!(Nothing, input, output_deriv, results, 1) else a, b = input p = 0 if istracked(a) p += 1 - diffresult_increment_deriv!(a, output_deriv, results, p) + diffresult_increment_deriv!(Nothing, a, output_deriv, results, p) end if istracked(b) p += 1 - diffresult_increment_deriv!(b, output_deriv, results, p) - end - end - unseed!(output) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(broadcast)}) - input = instruction.input - output = instruction.output - output_deriv = deriv(output) - results, _, a_bound, b_bound = instruction.cache - if istracked(input) - if size(input) == size(output_deriv) - diffresult_increment_deriv!(input, output_deriv, results, 1) - else - diffresult_increment_deriv!(input, output_deriv, results, 1, a_bound) - end - else - a, b = input - p = 0 - if size(a) == size(b) - if istracked(a) - p += 1 - diffresult_increment_deriv!(a, output_deriv, results, p) - end - if istracked(b) - p += 1 - diffresult_increment_deriv!(b, output_deriv, results, p) - end - else - if istracked(a) - p += 1 - diffresult_increment_deriv!(a, output_deriv, results, p, a_bound) - end - if istracked(b) - p += 1 - diffresult_increment_deriv!(b, output_deriv, results, p, b_bound) - end - end - end - unseed!(output) - return nothing -end - -############################# -# built-in infix operations # -############################# - -const TrackedType = Union{TrackedArray,TrackedReal} - -# dispatch # -#----------# - -for (F, broadcast_f) in ((typeof(+), :broadcast_plus), - (typeof(-), :broadcast_minus), - (typeof(*), :broadcast_mul), - (typeof(/), :broadcast_rdiv), - (typeof(\), :broadcast_ldiv), - (typeof(^), :broadcast_pow)) - @eval begin - @inline Base.broadcast(::$F, x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {X,Y,D} = $(broadcast_f)(x, y, D) - @inline Base.broadcast(::$F, x::TrackedReal{X,D}, y::TrackedArray{Y,D}) where {X,Y,D} = $(broadcast_f)(x, y, D) - @inline Base.broadcast(::$F, x::TrackedArray{X,D}, y::TrackedReal{Y,D}) where {X,Y,D} = $(broadcast_f)(x, y, D) - end - for A in ARRAY_TYPES - @eval begin - @inline Base.broadcast(::$F, x::TrackedArray{X,D}, y::$A{<:Real}) where {X,D} = $(broadcast_f)(x, y, D) - @inline Base.broadcast(::$F, x::$A{<:Real}, y::TrackedArray{Y,D}) where {Y,D} = $(broadcast_f)(x, y, D) - @inline Base.broadcast(::$F, x::TrackedReal{X,D}, y::$A{<:Real}) where {X,D} = $(broadcast_f)(x, y, D) - @inline Base.broadcast(::$F, x::$A{<:Real}, y::TrackedReal{Y,D}) where {Y,D} = $(broadcast_f)(x, y, D) - end - end - for R in REAL_TYPES - @eval begin - @inline Base.broadcast(::$F, x::TrackedArray{X,D}, y::$R) where {X,D} = $(broadcast_f)(x, y, D) - @inline Base.broadcast(::$F, x::$R, y::TrackedArray{Y,D}) where {Y,D} = $(broadcast_f)(x, y, D) + diffresult_increment_deriv!(Nothing, b, output_deriv, results, p) end end -end - -# .+ # -#----# - -function broadcast_plus(x, y, ::Type{D}) where D - tp = tape(x, y) - out = track(value(x) .+ value(y), D, tp) - cache = (index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, (broadcast, +), (x, y), out, cache) - return out -end - -@noinline function special_forward_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(+)}}) - a, b = instruction.input - output = instruction.output - pull_value!(a) - pull_value!(b) - broadcast!(+, value(output), value(a), value(b)) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(+)}}) - a, b = instruction.input - output = instruction.output - output_deriv = deriv(output) - a_bound, b_bound = instruction.cache - istracked(a) && broadcast_increment_deriv!(a, output_deriv, a_bound) - istracked(b) && broadcast_increment_deriv!(b, output_deriv, b_bound) - unseed!(output) - return nothing -end - -# .- # -#----# - -function broadcast_minus(x, y, ::Type{D}) where D - tp = tape(x, y) - out = track(value(x) .- value(y), D, tp) - cache = (index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, (broadcast, -), (x, y), out, cache) - return out -end - -@noinline function special_forward_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(-)}}) - a, b = instruction.input - output = instruction.output - pull_value!(a) - pull_value!(b) - broadcast!(-, value(output), value(a), value(b)) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(-)}}) - a, b = instruction.input - output = instruction.output - output_deriv = deriv(output) - a_bound, b_bound = instruction.cache - istracked(a) && broadcast_increment_deriv!(a, output_deriv, a_bound) - istracked(b) && broadcast_decrement_deriv!(b, output_deriv, b_bound) - unseed!(output) - return nothing -end - -# .* # -#----# - -function broadcast_mul(x, y, ::Type{D}) where D - tp = tape(x, y) - out = track(value(x) .* value(y), D, tp) - cache = (index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, (broadcast, *), (x, y), out, cache) - return out -end - -@noinline function special_forward_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(*)}}) - a, b = instruction.input - output = instruction.output - pull_value!(a) - pull_value!(b) - broadcast!(*, value(output), value(a), value(b)) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(*)}}) - a, b = instruction.input - output = instruction.output - output_deriv = deriv(output) - a_bound, b_bound = instruction.cache - istracked(a) && broadcast_increment_deriv!(a, output_deriv, value(b), a_bound, b_bound) - istracked(b) && broadcast_increment_deriv!(b, output_deriv, value(a), b_bound, a_bound) - unseed!(output) - return nothing -end - -# ./ # -#----# - -denom_partials_kernel(n::Real, d::Real) = -(n / d) / d -denom_partials(n, d) = broadcast(denom_partials_kernel, n, d) -denom_partials!(out::AbstractArray, n, d) = (broadcast!(denom_partials_kernel, out, n, d); nothing) - -function broadcast_rdiv(x, y, ::Type{D}) where D - tp = tape(x, y) - out = track(value(x) ./ value(y), D, tp) - d_partials = denom_partials(value(x), value(y)) - cache = (d_partials, index_bound(x, out), index_bound(y, out), - index_bound(d_partials, out)) - record!(tp, SpecialInstruction, (broadcast, /), (x, y), out, cache) - return out -end - -@noinline function special_forward_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(/)}}) - a, b = instruction.input - a_value, b_value = value(a), value(b) - d_partials = first(instruction.cache) - output = instruction.output - pull_value!(a) - pull_value!(b) - broadcast!(/, value(output), a_value, b_value) - istracked(b) && denom_partials!(d_partials, a_value, b_value) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(/)}}) - a, b = instruction.input - output = instruction.output - output_deriv = deriv(output) - d_partials, a_bound, b_bound, d_partials_bound = instruction.cache - istracked(a) && broadcast_increment_div_deriv!(a, output_deriv, value(b), a_bound, b_bound) - istracked(b) && broadcast_increment_deriv!(b, output_deriv, d_partials, b_bound, d_partials_bound) - unseed!(output) - return nothing -end - -# .\ # -#----# - -function broadcast_ldiv(x, y, ::Type{D}) where D - tp = tape(x, y) - out = track(value(x) .\ value(y), D, tp) - d_partials = denom_partials(value(y), value(x)) - cache = (d_partials, index_bound(x, out), index_bound(y, out), - index_bound(d_partials, out)) - record!(tp, SpecialInstruction, (broadcast, \), (x, y), out, cache) - return out -end - -@noinline function special_forward_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(\)}}) - a, b = instruction.input - a_value, b_value = value(a), value(b) - d_partials = first(instruction.cache) - output = instruction.output - pull_value!(a) - pull_value!(b) - broadcast!(\, value(output), a_value, b_value) - istracked(a) && denom_partials!(d_partials, b_value, a_value) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(\)}}) - a, b = instruction.input - output = instruction.output - output_deriv = deriv(output) - d_partials, a_bound, b_bound, d_partials_bound = instruction.cache - istracked(a) && broadcast_increment_deriv!(a, output_deriv, d_partials, a_bound, d_partials_bound) - istracked(b) && broadcast_increment_div_deriv!(b, output_deriv, value(a), b_bound, a_bound) - unseed!(output) - return nothing -end - -# .^ # -#----# - -base_partials_kernel(b::Real, e::Real) = e * b^(e - 1) -base_partials(b, e) = broadcast(base_partials_kernel, b, e) -base_partials!(out::AbstractArray, b, e) = (broadcast!(base_partials_kernel, out, b, e); nothing) - -exp_partials_kernel(b::Real, e::Real) = log(b) * b^e -exp_partials(b, e) = broadcast(exp_partials_kernel, b, e) -exp_partials!(out::AbstractArray, b, e) = (broadcast!(exp_partials_kernel, out, b, e); nothing) - -function pow_cache(x, y) - pow_x = istracked(x) ? base_partials(value(x), value(y)) : value(x) - pow_y = istracked(y) ? exp_partials(value(x), value(y)) : value(y) - return (pow_x, pow_y) -end - -function broadcast_pow(x, y, ::Type{D}) where D - tp = tape(x, y) - out = track(value(x) .^ value(y), D, tp) - bs_partials, ex_partials = pow_cache(x, y) - cache = (bs_partials, ex_partials, - index_bound(x, out), index_bound(y, out), - index_bound(bs_partials, out), index_bound(ex_partials, out)) - record!(tp, SpecialInstruction, (broadcast, ^), (x, y), out, cache) - return out -end - -@noinline function special_forward_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(^)}}) - a, b = instruction.input - a_value, b_value = value(a), value(b) - bs_partials, ex_partials = instruction.cache - output = instruction.output - pull_value!(a) - pull_value!(b) - broadcast!(^, value(output), a_value, b_value) - istracked(a) && base_partials!(bs_partials, a_value, b_value) - istracked(b) && exp_partials!(ex_partials, a_value, b_value) - return nothing -end - -@noinline function special_reverse_exec!(instruction::SpecialInstruction{Tuple{typeof(broadcast),typeof(^)}}) - a, b = instruction.input - output = instruction.output - output_deriv = deriv(output) - bs_partials, ex_partials, a_bound, b_bound, - bs_partials_bound, ex_partials_bound = instruction.cache - istracked(a) && broadcast_increment_deriv!(a, output_deriv, bs_partials, a_bound, bs_partials_bound) - istracked(b) && broadcast_increment_deriv!(b, output_deriv, ex_partials, b_bound, ex_partials_bound) unseed!(output) return nothing end diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index 570a18b0..1ab669ef 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -6,7 +6,7 @@ amount scaled by the output's derivative(s). Sometimes, extra partial informatio form of normal scalars/arrays, or `Dual` numbers) needs to accounted for as well. Often, a function's input and output values are not similarly shaped. To account for these -cases, `broadcast` and `reduce` versions of the propagation functions have been implemented. +cases, `diffresult` and `reduce` versions of the propagation functions have been implemented. Accumulations into a scalar derivative go through `sum`, whose pairwise reduction is both more accurate and faster than a sequential loop. @@ -96,122 +96,75 @@ end # diffresult_increment_deriv! # ############################### -@inline getpartial(r::DiffResults.ImmutableDiffResult{1,V,Tuple{D}}, p) where {V,D<:AbstractArray} = DiffResults.derivative(r)[p] -@inline getpartial(r::DiffResults.ImmutableDiffResult{1,V,Tuple{D}}, p) where {V,D<:Number} = DiffResults.derivative(r) +@inline getpartial(::Type, r::DiffResults.ImmutableDiffResult{1,V,Tuple{D}}, p) where {V,D<:AbstractArray} = DiffResults.derivative(r)[p] +@inline getpartial(::Type, r::DiffResults.ImmutableDiffResult{1,V,Tuple{D}}, p) where {V,D<:Number} = DiffResults.derivative(r) +@inline getpartial(::Type{T}, d::ForwardDiff.Dual, p) where {T} = ForwardDiff.partials(T, d, p) +@inline getpartial(::Type, x::Real, p) = zero(x) -function diffresult_increment_deriv!(input::AbstractArray, x::AbstractArray, - results, p::Int) - for i in eachindex(x, results) - increment_deriv!(input, x[i] * getpartial(results[i], p), i) - end - return nothing +# a partial known in closed form: the argument collects `op(seed, args...)` per element, with +# `args` naming broadcast arguments. `op` meets the seed, so `/` forms no reciprocal. +struct Contract{Op,A<:Tuple} + op::Op + args::A end -function diffresult_increment_deriv!(input::AbstractArray, x::AbstractArray, - results, p::Int, bound::CartesianIndex) - axes(x) == axes(results) || - throw(DimensionMismatch("`x` and `results` must have the same indices")) - for (xi, r, i) in zip(x, results, CartesianIndices(size(x))) - increment_deriv!(input, xi * getpartial(r, p), min(bound, i)) - end - return nothing -end +_at(::Val{j}, i, args) where {j} = _elem(value(args[j]), i) +_elem(v::Real, i) = v +_elem(v::AbstractArray, i) = v[i] -function diffresult_increment_deriv!(input::TrackedReal, x::AbstractArray, - results, p::Int, ::Nothing) - inds = eachindex(x, results) - isempty(inds) && return nothing - pull_deriv!(input) - input.deriv += sum(i -> x[i] * getpartial(results[i], p), inds) - push_deriv!(input) - return nothing -end +_contract(e::Contract, seed, i, args) = e.op(seed, map(a -> _at(a, i, args), e.args)...) -############################## -# broadcast_increment_deriv! # -############################## - -# without partials # -#------------------# - -function broadcast_increment_deriv!(input::AbstractArray, x::AbstractArray, - bound::CartesianIndex) - for (xi, i) in zip(x, CartesianIndices(size(x))) - increment_deriv!(input, xi, min(bound, i)) +function diffresult_increment_deriv!(::Type{T}, input::AbstractArray, x::AbstractArray, + results::AbstractArray, p::Int) where {T} + for i in eachindex(input, x, results) + increment_deriv!(input, x[i] * getpartial(T, results[i], p), i) end return nothing end -function broadcast_increment_deriv!(input::TrackedReal, x::AbstractArray, ::Nothing) - isempty(x) && return nothing - pull_deriv!(input) - input.deriv += sum(x) - push_deriv!(input) - return nothing -end - -# with partials # -#---------------# - -@inline broadcast_increment_deriv!(input, x, partials, input_bound, partials_bound) = - _broadcast_increment_deriv!(*, input, x, partials, input_bound, partials_bound) - -@inline broadcast_increment_div_deriv!(input, x, partials, input_bound, partials_bound) = - _broadcast_increment_deriv!(/, input, x, partials, input_bound, partials_bound) - -# with partial array # -#--------------------# - -function _broadcast_increment_deriv!(op::F, input::AbstractArray, x::AbstractArray, - partials::AbstractArray, - input_bound::CartesianIndex, - partials_bound::CartesianIndex) where {F} - for (xi, i) in zip(x, CartesianIndices(size(x))) - current_deriv = op(xi, partials[min(partials_bound, i)]) - increment_deriv!(input, current_deriv, min(input_bound, i)) +function diffresult_increment_deriv!(::Type{T}, input::AbstractArray, x::AbstractArray, + results::AbstractArray, p::Int, + bound::CartesianIndex) where {T} + axes(x) == axes(results) || + throw(DimensionMismatch("`x` and `results` must have the same indices")) + for (xi, r, i) in zip(x, results, CartesianIndices(size(x))) + increment_deriv!(input, xi * getpartial(T, r, p), min(bound, i)) end return nothing end -function _broadcast_increment_deriv!(op::F, input::TrackedReal, x::AbstractArray, - partials::AbstractArray, ::Nothing, - ::CartesianIndex) where {F} - inds = eachindex(x, partials) +function diffresult_increment_deriv!(::Type{T}, input::TrackedReal, x::AbstractArray, + results::AbstractArray, p::Int, ::Nothing) where {T} + inds = eachindex(x, results) isempty(inds) && return nothing pull_deriv!(input) - input.deriv += sum(i -> op(x[i], partials[i]), inds) + input.deriv += sum(i -> x[i] * getpartial(T, results[i], p), inds) push_deriv!(input) return nothing end -# with partial scalar # -#---------------------# - -function _broadcast_increment_deriv!(op::F, input::AbstractArray, x::AbstractArray, - partial::Real, input_bound::CartesianIndex, - ::Nothing) where {F} - for (xi, i) in zip(x, CartesianIndices(size(x))) - increment_deriv!(input, op(xi, partial), min(input_bound, i)) +# a closed-form partial is named rather than stored, so `args` stands where an index would +function diffresult_increment_deriv!(::Type, input::AbstractArray, x::AbstractArray, + results::Contract, args::Tuple) + for i in eachindex(input, x) + increment_deriv!(input, _contract(results, x[i], i, args), i) end return nothing end -############################## -# broadcast_decrement_deriv! # -############################## - -function broadcast_decrement_deriv!(input::AbstractArray, x::AbstractArray, - bound::CartesianIndex) - for (xi, i) in zip(x, CartesianIndices(size(x))) - decrement_deriv!(input, xi, min(bound, i)) +function diffresult_increment_deriv!(::Type, input::AbstractArray, x::AbstractArray, + results::Contract, args::Tuple, bound::CartesianIndex) + for i in CartesianIndices(size(x)) + increment_deriv!(input, _contract(results, x[i], i, args), min(bound, i)) end return nothing end -function broadcast_decrement_deriv!(input::TrackedReal, x::AbstractArray, ::Nothing) +function diffresult_increment_deriv!(::Type, input::TrackedReal, x::AbstractArray, + results::Contract, args::Tuple, ::Nothing) isempty(x) && return nothing pull_deriv!(input) - input.deriv -= sum(x) + input.deriv += sum(i -> _contract(results, x[i], i, args), eachindex(x)) push_deriv!(input) return nothing end diff --git a/src/tracked.jl b/src/tracked.jl index 50c77c4e..e0486947 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -202,7 +202,7 @@ push_deriv!(x::AbstractArray) = (istracked(x) && foreach(push_deriv!, x); nothin # seed/unseed # #-------------# -seed!(x) = nothing +seed!(::Real) = nothing seed!(t::TrackedReal) = (t.deriv = one(derivtype(t)); push_deriv!(t); nothing) seed!(t::TrackedArray, i) = (t.deriv[i] = one(derivtype(t)); nothing) seed!(x::AbstractArray, i) = seed!(x[i]) @@ -517,6 +517,10 @@ track(x::Real, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = T track(x::AbstractArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = TrackedArray(x, fill!(similar(x, D), zero(D)), tp) +# every forward pass writes into the value buffer, which a `TrackedArray` rejects +# TODO: self-nesting, and the `value`/`track` unwrapping elsewhere, risk perturbation confusion +track(x::TrackedArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = track(collect(x), D, tp) + track!(t::TrackedArray, x::AbstractArray) = (value!(t, x); unseed!(t); t) track!(t::TrackedReal, x::Real) = (value!(t, x); unseed!(t); t) diff --git a/test/PrecisionTests.jl b/test/PrecisionTests.jl index 29d777aa..3da81cd1 100644 --- a/test/PrecisionTests.jl +++ b/test/PrecisionTests.jl @@ -31,10 +31,9 @@ setprecision(BigFloat, 256) do @test ReverseDiff.gradient(f, x) ≈ ∇f(x) rtol=rtol end - # fused broadcasts lose precision in `ForwardDiff._div_partials` (upstream) let f = y -> sum(y .^ 2 ./ 7) test_println("BigFloat gradients", f) - @test_broken ReverseDiff.gradient(f, x) ≈ 2 .* x ./ 7 rtol=rtol + @test ReverseDiff.gradient(f, x) ≈ 2 .* x ./ 7 rtol=rtol end end diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index 0ef5648a..f033d9cb 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -389,14 +389,14 @@ ta = TrackedArray(varr, darr, InstructionTape()) tr = TrackedReal(v, d) trs = Any[tr, ta[1], ta[2]] -@test ReverseDiff.seed!(nothing) === nothing -@test ReverseDiff.seed!([1,2,3]) === nothing +@test_throws MethodError ReverseDiff.seed!(nothing) +@test_throws MethodError ReverseDiff.seed!([1,2,3]) @test ReverseDiff.seed!(1.0) === nothing @test ReverseDiff.unseed!(nothing) === nothing @test ReverseDiff.unseed!([1,2,3]) === nothing @test ReverseDiff.unseed!(1.0) === nothing -ReverseDiff.seed!(ta) +@test_throws MethodError ReverseDiff.seed!(ta) @test ReverseDiff.value(ta) === varr == varr_copy @test ReverseDiff.deriv(ta) === darr == darr_copy ReverseDiff.unseed!(ta) @@ -416,7 +416,7 @@ ReverseDiff.deriv!(tr, d) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === d -ReverseDiff.seed!(trs) +@test_throws MethodError ReverseDiff.seed!(trs) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === d @test ReverseDiff.value(ta) === varr == varr_copy @@ -445,7 +445,7 @@ ReverseDiff.pull_deriv!(trs) @test ReverseDiff.deriv(trs[2]) === darr[1] @test ReverseDiff.deriv(trs[3]) === darr[2] -ReverseDiff.seed!((tr, ta, trs)) +@test_throws MethodError ReverseDiff.seed!((tr, ta, trs)) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === d @test ReverseDiff.value(ta) === varr == varr_copy @@ -780,7 +780,8 @@ end ta_nested = ReverseDiff.track(collect(ta), eltype(ta), tp) @test float(ta_nested) === ta_nested -@test all(samefields.(ta, copyto!(similar(ta), ta))) +# mapped, not broadcast: `samefields` compares tracked fields, which `Dual`s do not carry +@test all(map(samefields, ta, copyto!(similar(ta), ta))) #################### # `Real` Interface # diff --git a/test/api/JacobianTests.jl b/test/api/JacobianTests.jl index 0f26db34..39cc9b2e 100644 --- a/test/api/JacobianTests.jl +++ b/test/api/JacobianTests.jl @@ -333,4 +333,9 @@ f251(x) = x .^ 2 end end +@testset "outputs whose element type is not `Real`" begin + @test ReverseDiff.jacobian(x -> Any[2 * x[1], 3 * x[2]], [1.0, 2.0]) == [2.0 0.0; 0.0 3.0] + @test ReverseDiff.jacobian(x -> [1.0, 2.0], rand(3)) == zeros(2, 3) +end + end # module diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl new file mode 100644 index 00000000..c9dc2f2d --- /dev/null +++ b/test/derivatives/BroadcastTests.jl @@ -0,0 +1,360 @@ +module BroadcastTests + +using ReverseDiff + +using ForwardDiff +using LinearAlgebra +using SpecialFunctions +using StaticArrays +using Test + +include("../utils.jl") + +using Base.Broadcast: AbstractArrayStyle, BroadcastStyle, Broadcasted, DefaultArrayStyle, + Style, Unknown, result_style +using ReverseDiff: TrackedArray, TrackedReal, TrackedStyle + +# stands in for a foreign array style such as `CuArrayStyle` or `StructuredMatrixStyle` +struct ForeignStyle <: AbstractArrayStyle{Any} end +ForeignStyle(::Val) = ForeignStyle() + +# a foreign 0-dimensional array, whose style carries `Any` rather than a dimension +struct Foreign0d{T} <: AbstractArray{T,0} + x::T +end +Base.size(::Foreign0d) = () +Base.getindex(f::Foreign0d) = f.x +Base.getindex(f::Foreign0d, ::CartesianIndex{0}) = f.x +Base.Broadcast.BroadcastStyle(::Type{<:Foreign0d}) = ForeignStyle() +Base.similar(::Broadcasted{ForeignStyle}, ::Type{T}, axs) where {T} = similar(Array{T}, axs) + +############################################################################################ + +@testset "`BroadcastStyle` tracks dimensionality" begin + tp = InstructionTape() + + @test BroadcastStyle(typeof(track(rand(3), tp))) === TrackedStyle{1}() + @test BroadcastStyle(typeof(track(rand(3, 3), tp))) === TrackedStyle{2}() + # `Base` styles scalars as `AbstractArrayStyle{0}` + @test BroadcastStyle(typeof(track(rand(), tp))) === TrackedStyle{0}() + @test TrackedStyle{1}(Val(2)) === TrackedStyle{2}() +end + +@testset "`BroadcastStyle` precedence" begin + for (a, b, expected) in ((TrackedStyle{0}(), DefaultArrayStyle{0}(), TrackedStyle{0}()), + (TrackedStyle{0}(), DefaultArrayStyle{2}(), TrackedStyle{2}()), + (TrackedStyle{1}(), DefaultArrayStyle{1}(), TrackedStyle{1}()), + (TrackedStyle{1}(), TrackedStyle{2}(), TrackedStyle{2}()), + (TrackedStyle{1}(), Unknown(), TrackedStyle{1}()), + (TrackedStyle{1}(), ForeignStyle(), TrackedStyle{Any}()), + # a scalar loses to a tuple, as in `Base` + (TrackedStyle{0}(), Style{Tuple}(), Style{Tuple}()), + (TrackedStyle{1}(), Style{Tuple}(), TrackedStyle{1}())) + @test result_style(a, b) === expected + @test result_style(b, a) === expected + end +end + +@testset "foreign array styles keep their values tracked" begin + a, v = rand(3, 3), rand(3) + d = Diagonal(rand(3)) + u = UpperTriangular(rand(3, 3)) + s = SVector(1.0, 2.0, 3.0) + + @test result_style(TrackedStyle{2}(), BroadcastStyle(typeof(d))) === TrackedStyle{2}() + @test result_style(TrackedStyle{2}(), BroadcastStyle(typeof(u))) === TrackedStyle{2}() + @test result_style(TrackedStyle{1}(), BroadcastStyle(typeof(s))) === TrackedStyle{1}() + + tp = InstructionTape() + @test track(a, tp) .* d isa TrackedArray + @test track(v, tp) .* s isa TrackedArray + # a static container survives into the tracked value + @test value(track(v, tp) .* s) isa StaticArray + + @test ReverseDiff.gradient(x -> sum(x .* d), a) ≈ Matrix(d) + @test ReverseDiff.gradient(x -> sum(x .* u), a) ≈ Matrix(u) + @test ReverseDiff.gradient(x -> sum(x .* s), v) ≈ s +end + +@testset "no `BroadcastStyle` ambiguities" begin + # `ReverseDiff` has unrelated pre-existing ambiguities, so filter to broadcasting + @test isempty(filter(Test.detect_ambiguities(ReverseDiff)) do (m1, m2) + m1.name === :BroadcastStyle || m2.name === :BroadcastStyle + end) +end + +@testset "scalar broadcasting matches `Base`" begin + tr = track(2.0, InstructionTape()) + + @test exp.(tr) isa TrackedReal + @test tr .+ tr isa TrackedReal + @test tr .+ [1.0, 2.0, 3.0] isa TrackedArray + @test tr .+ (1, 2, 3) isa Tuple +end + +@testset "`f.(x)` and `broadcast(f, x)` agree" begin + a, b = rand(3), rand(3) + + for (dotted, called) in ((tp -> exp.(track(a, tp)), tp -> broadcast(exp, track(a, tp))), + (tp -> track(a, tp) .+ track(b, tp), tp -> broadcast(+, track(a, tp), track(b, tp))), + (tp -> atan.(track(a, tp), b), tp -> broadcast(atan, track(a, tp), b))) + tp1, tp2 = InstructionTape(), InstructionTape() + x, y = dotted(tp1), called(tp2) + + @test value(x) == value(y) + @test length(tp1) == length(tp2) == 1 + @test tp1[1].func === tp2[1].func + end +end + +@testset "`@forward` wrappers do not force the fallback path" begin + tp = InstructionTape() + x = track(rand(3, 3), tp) + + y = broadcast(ReverseDiff.ForwardOptimize(exp), x) + + @test y isa TrackedArray + @test length(tp) == 1 +end + +@testset "scalar-array broadcasting preserves the derivative type" begin + # `Base` routes scalar-times-array through broadcasting + a = rand(1, 1) + f = x -> sum(first(x) * (x * x) + x) # == x[1]^3 + x[1] + + @test ReverseDiff.hessian(f, a) ≈ [6 * a[1];;] +end + +@testset "broadcast expressions stay fused" begin + tp = InstructionTape() + x, y = track(rand(3), tp), track(rand(3), tp) + + sum((x .+ y) .* sin.(x)) + + # one instruction for the fused expression, one for `sum` + @test length(tp) == 2 +end + +@testset "`TrackedArray`s are rejected as broadcast destinations" begin + a = rand(4) + msg = "`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead." + + tp = InstructionTape() + x = track(copy(a), tp) + @test_throws ArgumentError(msg) track(zeros(4), tp) .= exp.(x) + @test_throws ArgumentError(msg) track(zeros(4), tp) .= a + # nothing may be recorded before the failure + @test isempty(tp) + + # an untracked container of tracked elements is a valid destination + dest = Vector{TrackedReal{Float64,Float64,Nothing}}(undef, 4) + dest .= exp.(x) + @test value.(dest) ≈ exp.(a) + # the elements carry the tape, so it must not be dropped on the way in + @test all(d -> tape(d) === tp, dest) +end + +@testset "an untracked argument's `NaN` partial stays out of the tracked partials" begin + # `DiffRules` gives `besselj` a `NaN` partial for its order, which is untracked here + nu, z = [0.5, 1.5, 2.5], 0.7 + # J_ν'(z) = (J_{ν-1}(z) - J_{ν+1}(z)) / 2 + expected = sum((besselj(n - 1, z) - besselj(n + 1, z)) / 2 for n in nu) + + @test ReverseDiff.gradient(y -> sum(besselj.(nu, y)), [z]) ≈ [expected] +end + +@testset "an array of `TrackedReal`s is differentiated like a `TrackedArray`" begin + a, b = rand(3), rand(3) + tp = InstructionTape() + # a container of tracked elements, which `track` itself never produces + x = map(xi -> track(xi, tp), a) + + y = x .* b + + @test y isa TrackedArray + @test value(y) ≈ a .* b + @test length(tp) == 1 + @test tp[1].func === ReverseDiff.∇broadcast + + ReverseDiff.seed!(sum(y)) + ReverseDiff.reverse_pass!(tp) + # `deriv.(x)` would be traced like any other broadcast, as it is for a `TrackedArray` + @test map(deriv, x) ≈ b +end + +@testset "functions constant in their argument" begin + a = rand(3) + + @test ReverseDiff.gradient(x -> sum((t -> 1.0).(x)), a) == zeros(3) + @test ReverseDiff.gradient(x -> sum(oneunit.(x)), a) == zeros(3) +end + +@testset "a type-unstable function keeps its partials" begin + # the integer literal makes `Base` widen the results to an abstract element type + relu(t) = t > 0 ? t : 0 + a = [-1.0, 2.0] + + tp = InstructionTape() + y = relu.(track(copy(a), tp)) + + @test y isa TrackedArray + @test value(y) == [0.0, 2.0] + @test ReverseDiff.gradient(x -> sum(relu.(x)), a) == [0.0, 1.0] +end + +@testset "a foreign tag is never read as our own" begin + tagA = typeof(ForwardDiff.Tag(sin, Float64)) + tagB = typeof(ForwardDiff.Tag(cos, Float64)) + # `Tag`s are ordered by first use, so pin the order before relying on it + @test ForwardDiff.:≺(tagA, tagB) + + concrete = ForwardDiff.Dual{tagA,Float64,1}[ForwardDiff.Dual{tagA}(2.0, 3.0)] + widened = Union{Float64,ForwardDiff.Dual{tagA,Float64,1}}[ + 1.0, ForwardDiff.Dual{tagA}(2.0, 3.0)] + nested = Union{Float64,ForwardDiff.Dual{tagB,Float64,1}}[1.0] + + @test ReverseDiff.trackresults(tagB, concrete, identity, identity, (), Float64) === concrete + @test ReverseDiff.trackresults(tagB, widened, identity, identity, (), Float64) === widened + @test_throws ForwardDiff.DualMismatchError ReverseDiff.trackresults( + tagA, nested, identity, identity, (), Float64) + + @test ReverseDiff.getpartial(tagA, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 2.0 + @test ReverseDiff.getpartial(tagA, 1.0, 1) == 0.0 + @test ReverseDiff.getpartial(tagB, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 0.0 + @test_throws ForwardDiff.DualMismatchError ReverseDiff.getpartial( + tagA, ForwardDiff.Dual{tagB}(1.0, 2.0), 1) +end + +@testset "zero-dimensional arrays (#265)" begin + a = rand(1) + + @test ReverseDiff.gradient(x -> exp.(reshape(vec(x), ())), a) ≈ exp.(a) + + x0 = reshape(vec(track(copy(a), InstructionTape())), ()) + + # `Base` collapses a 0-d broadcast to a scalar but keeps `map` 0-dimensional + @test broadcast(exp, x0) isa TrackedReal + @test map(exp, x0) isa TrackedArray + @test ndims(map(exp, x0)) == 0 +end + +@testset "an untracked scalar argument survives nested differentiation (#214)" begin + a = [0.5, 1.5, 2.5] + b = [1.0, 2.0, 3.0] + c = 4.0 + # `map(*, x, x ./ (b .+ c))` sums to `sum(x .^ 2 ./ (b .+ c))` + expected = diagm(2 ./ (b .+ c)) + + @test ReverseDiff.hessian(x -> sum(map(*, x, x ./ (b .+ c))), a) ≈ expected + @test ReverseDiff.hessian(x -> sum(map(*, x, x ./ (b .+ [c]))), a) ≈ expected +end + +@testset "0-dimensional broadcasts whose style carries no dimension" begin + tp = InstructionTape() + tr = track(2.0, tp) + x0 = reshape(vec(track([3.0], tp)), ()) + + # `TrackedStyle{Any}` cannot report 0 dimensions, so only the axes mark these as scalar + @test result_style(TrackedStyle{0}(), ForeignStyle()) === TrackedStyle{Any}() + @test tr .+ Foreign0d(1.0) isa TrackedReal + @test x0 .+ Foreign0d(1.0) isa TrackedReal + + @test ReverseDiff.gradient(v -> sum(v .* Foreign0d(2.0)), [1.0, 2.0]) ≈ [2.0, 2.0] +end + +@testset "`NotTracked`" begin + f = ReverseDiff.NotTracked(t -> 2t) + + @test ReverseDiff.gradient(x -> sum(f.(x)), rand(4)) == fill(2.0, 4) +end + +@testset "a partial known in closed form stores nothing" begin + n = 10_000 + v, w = rand(n), rand(n) + + function recordbytes(f, args...) + tp = InstructionTape() + targs = map(a -> track(a, tp), args) + f(targs...) + return @allocated f(targs...) + end + + # the output's value and derivative take 8 B per element each, and one `Dual` per + # element would add at least 16 B more + for f in (t -> t .* 2.0, t -> 2.0 .* t, t -> t .+ 1.0, t -> 1.0 .- t, t -> t ./ 4.0, + t -> 4.0 .\ t, t -> identity.(t), t -> .-t) + @test recordbytes(f, v) < 24n + end + for f in ((a, b) -> a .- b, (a, b) -> a .* b) + @test recordbytes(f, v, w) < 24n + end + + # an untracked array argument is copied onto the tape, 8 B per element more + for f in (t -> t .+ w, t -> w .- t, t -> t .* w, t -> t ./ w, t -> w .\ t) + @test recordbytes(f, v) < 32n + end + + # a partial that depends on the point still needs one `Dual` per element + for f in (t -> exp.(t), t -> t .^ 2, t -> 2.0 ./ t) + @test recordbytes(f, v) > 24n + end + # a tracked denominator's partial, `-x/y^2`, is no argument of the broadcast + @test recordbytes(t -> w ./ t, v) > 32n +end + +@testset "a known partial differentiates as the general path does" begin + v, w = [1.0, 2.0, 3.0], [4.0, 5.0, 6.0] + + @test ReverseDiff.gradient(t -> sum(t .* 2.0), v) == fill(2.0, 3) + @test ReverseDiff.gradient(t -> sum(t .+ 1.0), v) == fill(1.0, 3) + @test ReverseDiff.gradient(t -> sum(1.0 .- t), v) == fill(-1.0, 3) + @test ReverseDiff.gradient(t -> sum(t ./ 4.0), v) == fill(0.25, 3) + @test ReverseDiff.gradient(t -> sum(4.0 .\ t), v) == fill(0.25, 3) + @test ReverseDiff.gradient(t -> sum(.-t), v) == fill(-1.0, 3) + @test ReverseDiff.gradient(t -> sum(t .+ w), v) == fill(1.0, 3) + @test ReverseDiff.gradient(t -> sum(t .* w), v) == w + @test ReverseDiff.gradient(t -> sum(t ./ w), v) == 1 ./ w + @test ReverseDiff.gradient(t -> sum(w .\ t), v) == 1 ./ w + @test ReverseDiff.gradient(t -> sum(t .* t), v) == 2 .* v + + # `-x/y^2` is no argument of the broadcast, so it is read off a `Dual` + @test ReverseDiff.gradient(t -> sum(w ./ t), v) == -w ./ v .^ 2 + + # an outer broadcast still clamps the reverse pass to the argument's own shape + @test ReverseDiff.gradient(t -> sum(t .+ w'), v) == fill(3.0, 3) + @test ReverseDiff.gradient(t -> sum(t .* w'), v) == fill(sum(w), 3) + @test ReverseDiff.gradient(t -> sum(t .* 2.0), Float64[]) == Float64[] + + # a tracked scalar collects the whole output instead of one element of it + @test ReverseDiff.gradient(t -> sum(t[1] .+ w), v) == [3.0, 0.0, 0.0] + @test ReverseDiff.gradient(t -> sum(t .* t[1]), v) == [2 * v[1] + sum(v[2:end]), 1.0, 1.0] + + # an argument `splitargs` holds back is not one a partial may name + @test ReverseDiff.gradient(t -> sum(broadcast(*, t, w, Ref(2.0))), v) == 2 .* w + + ga, gb = ReverseDiff.gradient((a, b) -> sum(a .- b), (v, w)) + @test ga == fill(1.0, 3) + @test gb == fill(-1.0, 3) +end + +@testset "a known partial does not overflow on a subnormal divisor" begin + # materializing `1/y` would give `Inf` before it ever meets the seed + @test ReverseDiff.gradient(y -> 1e-10 * sum(y ./ 1e-310), [1.0, 2.0]) == + fill(1e-10 / 1e-310, 2) +end + +@testset "a known partial replays its values" begin + doubled(t) = t .* 2.0 + tape = ReverseDiff.GradientTape(t -> sum(exp.(doubled(t))), [1.0, 2.0, 3.0]) + x = [0.5, 1.5, 2.5] + + # `exp`'s partial reads the doubled values, so a stale replay shows up in the gradient + @test ReverseDiff.gradient!(tape, x) ≈ 2 .* exp.(2 .* x) + + # an argument a partial names is read at replay time, not at record time + scaled(t) = t .* [2.0, 3.0, 4.0] + tape2 = ReverseDiff.GradientTape(t -> sum(scaled(t)), [1.0, 2.0, 3.0]) + @test ReverseDiff.gradient!(tape2, x) == [2.0, 3.0, 4.0] +end + +end # module diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index 7ce670a1..8a9deec4 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -52,7 +52,8 @@ function test_elementwise(f, fopt, x, tp) # record yt = broadcast(fopt, xt) @test yt == y - @test length(tp) == 1 + # a function constant in its argument has no derivative and is left off the tape + @test length(tp) == (yt isa ReverseDiff.TrackedArray ? 1 : 0) # reverse out = similar(y, (length(x), length(x))) diff --git a/test/runtests.jl b/test/runtests.jl index b53f0e80..f5d4bde4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -54,6 +54,12 @@ test_println(kind, f, pad=" ") = println(pad, "testing $(kind): `$(f)`...") println("done (took $t seconds).") end + @testset "BroadcastTests" begin + println("running BroadcastTests...") + t = @elapsed include("derivatives/BroadcastTests.jl") + println("done (took $t seconds).") + end + @testset "ArrayFunctionTests" begin println("running ArrayFunctionTests...") t = @elapsed include(joinpath(TESTDIR, "derivatives/ArrayFunctionTests.jl")) From 290b9d6f630e4d84defde7504aeda8f5d3273f62 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 22 Sep 2026 21:25:03 +0200 Subject: [PATCH 02/36] Recurse into the element type when looking for hidden tracked values MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `fieldcount(Vector{T})` is 0 on Julia 1.10 and 2 on 1.11, where `Array` gained `ref` and `size`, so asking it whether an array argument could hide a tracked value answered differently by version. On 1.10 that let `x[1] * [[x[2], x[3]]]` reach `∇broadcast`, which seeded `x[1]` and left the enclosing `Dual` to be contracted with the seed in the reverse pass, where a `Float64` derivative cannot hold it (#140). Ask the element type instead. An argument is seeded itself or elementwise, so its own tracked values are in plain sight; below it nothing is seeded, so recurse into what a container holds rather than ask about the container. `istypeorclosure` never answered whether its input was a type or a closure, and `mayhavetracked` is false for a `TrackedReal`, so both become `mayhidetracked`, which is what they decide. Also report an argument carrying a perturbation that the tape's derivative type cannot hold, rather than letting it fail as a `MethodError` from the reverse pass. Co-Authored-By: Claude Opus 5 (1M context) --- src/derivatives/broadcast.jl | 49 ++++++++++++++++++++---------- test/derivatives/BroadcastTests.jl | 13 ++++++++ 2 files changed, 46 insertions(+), 16 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index c456757f..7816b683 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -15,13 +15,23 @@ struct NotTracked{F} <: Function end (f::NotTracked{<:Union{Function, Type}})(args...; kwargs...) = f.f(args...; kwargs...) -istypeorclosure(::F) where {F} = _istypeorclosure(F) -istypeorclosure(::AbstractArray{F}) where {F} = _istypeorclosure(F) -istypeorclosure(::Base.RefValue{F}) where {F} = _istypeorclosure(F) -istypeorclosure(::AbstractArray{<:Real}) = false -istypeorclosure(::TrackedArray) = false -istypeorclosure(::Real) = false -@generated function _istypeorclosure(::Type{F}) where {F} +# an argument is seeded itself, or elementwise, so its own tracked values are in plain sight; a +# `Ref` is handed to `f` whole, so nothing inside one is +mayhidetracked(b::F) where {F} = _mayhidetracked(F) +mayhidetracked(::AbstractArray{F}) where {F} = _mayhidetracked(F) +mayhidetracked(::Base.RefValue{F}) where {F} = _mayhidetracked(F) +mayhidetracked(::AbstractArray{<:Real}) = false +mayhidetracked(::Real) = false +mayhidetracked(b::Type) = false +mayhidetracked(b::ForwardOptimize) = mayhidetracked(b.f) +mayhidetracked(b::SkipOptimize) = mayhidetracked(b.f) +mayhidetracked(b::Broadcasted) = mayhidetracked(b.f) || any(mayhidetracked, b.args) + +# below the argument nothing is seeded, so a tracked value is hidden wherever it sits: recurse +# into what a container holds rather than ask about the container +_mayhidetracked(::Type{<:NotTracked}) = false +_mayhidetracked(::Type{<:AbstractArray{F}}) where {F} = _mayhidetracked(F) +@generated function _mayhidetracked(::Type{F}) where {F} # `fieldcount` errors for types without a definite number of fields, such as # `Type{T}` and abstract types; be conservative in that case. hasfields = try @@ -32,14 +42,6 @@ istypeorclosure(::Real) = false return :($hasfields) end -mayhavetracked(b) = istypeorclosure(b) -mayhavetracked(b::Type) = false -mayhavetracked(b::NotTracked) = false -mayhavetracked(b::ForwardOptimize) = mayhavetracked(b.f) -mayhavetracked(b::SkipOptimize) = mayhavetracked(b.f) -mayhavetracked(b::Base.RefValue{<:NotTracked}) = false -mayhavetracked(b::Broadcasted) = mayhavetracked(b.f) || any(mayhavetracked, b.args) - struct TrackedStyle{N} <: AbstractArrayStyle{N} end (::Type{<:TrackedStyle})(::Val{N}) where {N} = TrackedStyle{N}() @@ -87,7 +89,7 @@ function Base.copy(_bc::Broadcasted{<:TrackedStyle}) f, args = flattened_bc.f, flattened_bc.args # only the arguments are seeded, so a tracked value reaching `f` by another route, such # as a closure capturing one, has to be traced scalar-wise - if mayhavetracked(_bc) + if mayhidetracked(_bc) axs = flattened_bc.axes style = typeof(Broadcast.combine_styles(map(recur_value, args)...)) return copy(Broadcasted{style, typeof(axs), typeof(f), typeof(args)}(f, args, axs)) @@ -257,8 +259,23 @@ replaycache(::Type{T}, results::AbstractArray, df, vf, targs) where {T} = replaycache(::Type, ::KnownPartials, df, vf, targs) = (vf, broadcast(vf, map(value, targs)...)) +# the seed is contracted with an argument, so a perturbation riding on one ends up in the +# derivative and `D` has to be able to hold it; a widened element type is a `Union` +@inline function checkargtags(::Type{D}, ::Type{E}) where {D, E} + if E isa Union + checkargtags(D, E.a) + checkargtags(D, E.b) + elseif ForwardDiff.tagtype(E) !== Nothing && !(promote_type(D, E) <: D) + throw(ArgumentError(LazyString("a broadcast argument with element type ", E, + " carries a perturbation that a derivative of type ", D, + " cannot hold"))) + end + return nothing +end + # `df` hands back `f`'s own result, so `results` keeps the element type `Base` would produce @inline function recordresults(::Type{T}, results, df, vf, targs, ::Type{D}) where {T, D} + foreach(t -> checkargtags(D, eltype(value(t))), targs) g, outvalue = replaycache(T, results, df, vf, targs) tp = tape(targs...) out = track(outvalue, D, tp) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index c9dc2f2d..2d536535 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -225,6 +225,19 @@ end tagA, ForwardDiff.Dual{tagB}(1.0, 2.0), 1) end +@testset "a perturbation the derivative cannot hold is rejected (#67, #168)" begin + # `a` reaches the broadcast as a `Dual`, but the tape's derivatives are `Float64` + g(a) = sum(ReverseDiff.gradient(x -> sum(x .* a), [1.0, 2.0])) + E = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g),Float64},Float64,1} + msg = "a broadcast argument with element type $E carries a perturbation that a derivative of type Float64 cannot hold" + + @test_throws ArgumentError(msg) ForwardDiff.derivative(g, 3.0) + + # a tape whose derivatives are themselves `Dual`s can hold it, so it is left alone + h(a) = ReverseDiff.gradient(x -> sum(x .* a), [ForwardDiff.Dual(1.0, 0.0)]) + @test h(3.0) == [ForwardDiff.Dual(3.0, 0.0)] +end + @testset "zero-dimensional arrays (#265)" begin a = rand(1) From d862c2b11cc8bf4af058632e2e57f18b541eb862 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 09:49:56 +0200 Subject: [PATCH 03/36] Cover the `TrackedStyle{Any}`, wrapper, `Type` and `Union` paths Co-Authored-By: Claude Opus 5.5 (1M context) --- test/derivatives/BroadcastTests.jl | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 2d536535..4c2e783f 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -47,6 +47,8 @@ end (TrackedStyle{1}(), TrackedStyle{2}(), TrackedStyle{2}()), (TrackedStyle{1}(), Unknown(), TrackedStyle{1}()), (TrackedStyle{1}(), ForeignStyle(), TrackedStyle{Any}()), + (TrackedStyle{Any}(), ForeignStyle(), TrackedStyle{Any}()), + (TrackedStyle{Any}(), DefaultArrayStyle{1}(), TrackedStyle{Any}()), # a scalar loses to a tuple, as in `Base` (TrackedStyle{0}(), Style{Tuple}(), Style{Tuple}()), (TrackedStyle{1}(), Style{Tuple}(), TrackedStyle{1}())) @@ -64,6 +66,7 @@ end @test result_style(TrackedStyle{2}(), BroadcastStyle(typeof(d))) === TrackedStyle{2}() @test result_style(TrackedStyle{2}(), BroadcastStyle(typeof(u))) === TrackedStyle{2}() @test result_style(TrackedStyle{1}(), BroadcastStyle(typeof(s))) === TrackedStyle{1}() + @test result_style(TrackedStyle{Any}(), BroadcastStyle(typeof(s))) === TrackedStyle{Any}() tp = InstructionTape() @test track(a, tp) .* d isa TrackedArray @@ -107,14 +110,28 @@ end end end -@testset "`@forward` wrappers do not force the fallback path" begin +@testset "`@forward` and `@skip` wrappers do not force the fallback path" begin + for wrapper in (ReverseDiff.ForwardOptimize, ReverseDiff.SkipOptimize) + tp = InstructionTape() + x = track(rand(3, 3), tp) + + y = broadcast(wrapper(exp), x) + + @test y isa TrackedArray + @test length(tp) == 1 + end +end + +@testset "a type broadcast as a function does not force the fallback path" begin tp = InstructionTape() - x = track(rand(3, 3), tp) + x = track(rand(3), tp) - y = broadcast(ReverseDiff.ForwardOptimize(exp), x) + y = Real.(x) @test y isa TrackedArray @test length(tp) == 1 + @test tp[1].func === ReverseDiff.∇broadcast + @test ReverseDiff.gradient(v -> sum(Real.(v)), [1.0, 2.0]) == [1.0, 1.0] end @testset "scalar-array broadcasting preserves the derivative type" begin @@ -233,6 +250,12 @@ end @test_throws ArgumentError(msg) ForwardDiff.derivative(g, 3.0) + # a widened element type is checked member by member + g2(a) = sum(ReverseDiff.gradient(x -> sum(x .* Union{Float64,typeof(a)}[1.0, a]), [1.0, 2.0])) + E2 = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g2),Float64},Float64,1} + msg2 = "a broadcast argument with element type $E2 carries a perturbation that a derivative of type Float64 cannot hold" + @test_throws ArgumentError(msg2) ForwardDiff.derivative(g2, 3.0) + # a tape whose derivatives are themselves `Dual`s can hold it, so it is left alone h(a) = ReverseDiff.gradient(x -> sum(x .* a), [ForwardDiff.Dual(1.0, 0.0)]) @test h(3.0) == [ForwardDiff.Dual(3.0, 0.0)] From 065b8780214607777d89863596629bf905eb4cb8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 10:56:01 +0200 Subject: [PATCH 04/36] Leave `map` over a `TrackedReal` to `Base` `map` iterates a scalar as a one-element collection, so the fused methods had nothing to fuse, and their reverse pass and replay both threw. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/elementwise.jl | 57 +++++----------------------- test/derivatives/ElementwiseTests.jl | 20 ++++++++++ 2 files changed, 30 insertions(+), 47 deletions(-) diff --git a/src/derivatives/elementwise.jl b/src/derivatives/elementwise.jl index b56122a3..8c3dfaf8 100644 --- a/src/derivatives/elementwise.jl +++ b/src/derivatives/elementwise.jl @@ -76,17 +76,11 @@ for g in (:map,), (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) @eval @inline Base.$(g)(f::typeof($M.$f), t::TrackedArray) = $(g)(ForwardOptimize(f), t) elseif arity == 2 (M, f) in SKIPPED_DIFFRULES && continue - @eval begin - @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) - @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::TrackedReal) = $(g)(ForwardOptimize(f), x, y) - @inline Base.$(g)(f::typeof($M.$f), x::TrackedReal, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) - end + @eval @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) for A in ARRAY_TYPES @eval begin @inline Base.$(g)(f::typeof($M.$f), x::$A, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::$A) = $(g)(ForwardOptimize(f), x, y) - @inline Base.$(g)(f::typeof($M.$f), x::$A, y::TrackedReal) = $(g)(ForwardOptimize(f), x, y) - @inline Base.$(g)(f::typeof($M.$f), x::TrackedReal, y::$A) = $(g)(ForwardOptimize(f), x, y) end end for R in REAL_TYPES @@ -114,34 +108,6 @@ for g in (:map,) return out end for A in ARRAY_TYPES - @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedReal{X,D}, y::$A) where {F,X,D} - result = DiffResults.DiffResult(zero(X), zero(D)) - df = let result=result - (vx, vy) -> let vy=vy - ForwardDiff.derivative!(result, s -> f.f(s, vy), vx) - end - end - results = $(g)(df, value(x), value(y)) - tp = tape(x) - out = track(map(DiffResults.value, results), D, tp) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, $(g), (x, y), out, cache) - return out - end - @eval function Base.$(g)(f::ForwardOptimize{F}, x::$A, y::TrackedReal{Y,D}) where {F,Y,D} - result = DiffResults.DiffResult(zero(Y), zero(D)) - df = let result=result - (vx, vy) -> let vx=vx - ForwardDiff.derivative!(result, s -> f.f(vx, s), vy) - end - end - results = $(g)(df, value(x), value(y)) - tp = tape(y) - out = track(map(DiffResults.value, results), D, tp) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, $(g), (x, y), out, cache) - return out - end @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::$A) where {F,X,D} result = DiffResults.GradientResult(SVector(zero(X))) df = (vx, vy) -> let vy=vy @@ -168,18 +134,15 @@ for g in (:map,) end end - for TX in (:TrackedArray, :TrackedReal), TY in (:TrackedArray, :TrackedReal) - TX == :TrackedReal && TY == :TrackedReal && continue - @eval function Base.$(g)(f::ForwardOptimize{F}, x::$(TX){X,D}, y::$(TY){Y,D}) where {F,X,Y,D} - result = DiffResults.GradientResult(SVector(zero(D), zero(D))) - df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) - results = $(g)(df, value(x), value(y)) - tp = tape(x, y) - out = track(map(DiffResults.value, results), D, tp) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, $(g), (x, y), out, cache) - return out - end + @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {F,X,Y,D} + result = DiffResults.GradientResult(SVector(zero(D), zero(D))) + df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) + results = $(g)(df, value(x), value(y)) + tp = tape(x, y) + out = track(map(DiffResults.value, results), D, tp) + cache = (results, df, index_bound(x, out), index_bound(y, out)) + record!(tp, SpecialInstruction, $(g), (x, y), out, cache) + return out end end diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index 8a9deec4..64268805 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -454,4 +454,24 @@ end @test ReverseDiff.gradient(v -> sum(v .* Base.RefValue{Any}(2.0)), [0.3, 0.5]) == [2.0, 2.0] end +@testset "`map` over a tracked scalar and an array" begin + c = [1.5, 2.5] + + # `map` stops at its shortest argument, a scalar counting as a one-element collection, + # so each case is `atan` of the first elements alone + cases = ((t -> sum(map(atan, t[1], c)), (a, b) -> [c[1], 0] / (a^2 + c[1]^2)), + (t -> sum(map(atan, c, t[1])), (a, b) -> [-c[1], 0] / (a^2 + c[1]^2)), + (t -> sum(map(atan, t, t[2])), (a, b) -> [b, -a] / (a^2 + b^2)), + (t -> sum(map(atan, t[2], t)), (a, b) -> [-b, a] / (a^2 + b^2))) + for (f, ∇f) in cases + v = [0.5, 0.7] + @test ReverseDiff.gradient(f, v) ≈ ∇f(v...) + + # the replay reaches the scalar through the tape + tape = ReverseDiff.GradientTape(f, v) + w = [0.2, 0.9] + @test ReverseDiff.gradient!(tape, w) ≈ ∇f(w...) + end +end + end # module From 5d3365ce01c7973f0d844adca4a5dde0ba94f9dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:05:03 +0200 Subject: [PATCH 05/36] Trace a broadcast with non-`Real` results scalar-wise A `TrackedArray` holds only `Real`s, so `Complex` or `Tuple` results were returned untracked and their partials silently dropped. The inferred element type now sends such broadcasts down the fallback before `f` is evaluated. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 17 +++++++++-------- test/derivatives/BroadcastTests.jl | 16 ++++++++++++++++ 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 7816b683..12976a33 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -69,7 +69,7 @@ Broadcast.BroadcastStyle(::TrackedStyle{Any}, ::DefaultArrayStyle) = TrackedStyl # the untracked style decides where the fallback re-dispatches, and a `CuArray` backing a # `TrackedArray` has to keep its own recur_value(xs) = xs -recur_value(xs::Union{TrackedReal, TrackedArray, AbstractArray{<:TrackedReal}}) = recur_value(value(xs)) +recur_value(xs::Union{TrackedReal, AbstractArray{<:TrackedReal}}) = recur_value(value(xs)) remove_not_tracked(f) = f remove_not_tracked(f::NotTracked) = f.f @@ -87,14 +87,15 @@ function Base.copy(_bc::Broadcasted{<:TrackedStyle}) bc = remove_not_tracked(_bc) flattened_bc = Base.Broadcast.flatten(bc) f, args = flattened_bc.f, flattened_bc.args - # only the arguments are seeded, so a tracked value reaching `f` by another route, such - # as a closure capturing one, has to be traced scalar-wise - if mayhidetracked(_bc) + vals = map(value, args) + # only the arguments are seeded, not e.g. a closure's captures, and a `TrackedArray` + # holds only `Real`s + if mayhidetracked(_bc) || !(Broadcast.combine_eltypes(f, vals) <: Real) axs = flattened_bc.axes - style = typeof(Broadcast.combine_styles(map(recur_value, args)...)) + style = typeof(Broadcast.combine_styles(map(recur_value, vals)...)) return copy(Broadcasted{style, typeof(axs), typeof(f), typeof(args)}(f, args, axs)) else - return ∇broadcast(f, args...) + return ∇broadcast(f, args, vals) end end @@ -234,11 +235,11 @@ broadcastresults(entries::Tuple, slots, df, vals) = # at least one argument has to be a non-0-dimensional array: `copy` sends the scalar and # 0-dimensional cases down `Base`'s `TrackedStyle{0}` route onto the scalar rules instead -@inline function ∇broadcast(f::F, args::Vararg{Any}) where {F} +@inline function ∇broadcast(f::F, args::Tuple, argvals::Tuple) where {F} inds, targs, untracked = splitargs(args) + _, vals, _ = splitargs(argvals) D = mapreduce(getouttype, promote_type, targs) slots, valP = trackedslots(targs) - vals = map(value, targs) # one tag for the whole broadcast keeps `results` concretely typed T = typeof(ForwardDiff.Tag(f, reduce(promote_type, map(dualvaltype, slots, vals)))) # `broadcast` calls `df` elementwise, so it receives one scalar per argument diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 4c2e783f..f32396b2 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -219,6 +219,22 @@ end @test ReverseDiff.gradient(x -> sum(relu.(x)), a) == [0.0, 1.0] end +@testset "a non-`Real` result keeps its partials" begin + a = [1.0, 2.0] + + # `real.(cis.(x))` would fuse into one broadcast with `Real` results, so the `Complex` + # and `Tuple` results are materialized in a broadcast of their own + @test ReverseDiff.gradient(x -> (y = cis.(x); sum(real.(y))), a) ≈ -sin.(a) + @test ReverseDiff.gradient(x -> (y = sincos.(x); sum(first.(y))), a) ≈ cos.(a) + @test ReverseDiff.gradient(x -> (y = broadcast(t -> (t, 2t), x); sum(last.(y))), a) == + [2.0, 2.0] + @test ReverseDiff.gradient(x -> (y = x .* cis.(x); sum(abs2.(y))), a) ≈ 2 .* a + @test ReverseDiff.jacobian(x -> (y = cis.(x); imag.(y)), a) ≈ diagm(cos.(a)) + + tape = ReverseDiff.GradientTape(x -> (y = cis.(x); sum(real.(y))), a) + @test ReverseDiff.gradient!(tape, [0.5, 1.5]) ≈ -sin.([0.5, 1.5]) +end + @testset "a foreign tag is never read as our own" begin tagA = typeof(ForwardDiff.Tag(sin, Float64)) tagB = typeof(ForwardDiff.Tag(cos, Float64)) From 1c336c91ee281aa56085f3238219d0779b647bb5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:21:45 +0200 Subject: [PATCH 06/36] Let `map` stop at its shortest argument again `map` over arrays of different lengths gives a shorter output, so the reverse pass iterates over the output and its partials, not the input. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/propagation.jl | 2 +- test/derivatives/ElementwiseTests.jl | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index 1ab669ef..052a24ce 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -116,7 +116,7 @@ _contract(e::Contract, seed, i, args) = e.op(seed, map(a -> _at(a, i, args), e.a function diffresult_increment_deriv!(::Type{T}, input::AbstractArray, x::AbstractArray, results::AbstractArray, p::Int) where {T} - for i in eachindex(input, x, results) + for i in eachindex(x, results) increment_deriv!(input, x[i] * getpartial(T, results[i], p), i) end return nothing diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index 64268805..53add506 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -474,4 +474,14 @@ end end end +@testset "`map` over arrays of different lengths" begin + v = [0.5, 0.7, 0.9] + c = [1.5, 2.5] + + for f in (t -> sum(map(*, t, c)), t -> sum(map(+, c, t)), + t -> sum(map(atan, t, 2 .* t[2:3]))) + @test ReverseDiff.gradient(f, v) ≈ ForwardDiff.gradient(f, v) + end +end + end # module From 46d20147e0dc098cb2bf1309d17052ba6b1254e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:29:59 +0200 Subject: [PATCH 07/36] Claim only arrays whose element type fixes the derivative type `getouttype` reads the derivative type off the element type, so a `Vector{TrackedReal}` left it `Union{}`. Such arrays now keep `Base`'s style and are traced with the scalar rules, as before. The `TrackedArray` methods are subsumed by the `AbstractArray` ones. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 5 ++--- test/derivatives/BroadcastTests.jl | 11 +++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 12976a33..477f6286 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -46,9 +46,9 @@ struct TrackedStyle{N} <: AbstractArrayStyle{N} end (::Type{<:TrackedStyle})(::Val{N}) where {N} = TrackedStyle{N}() -Broadcast.BroadcastStyle(::Type{<:TrackedArray{V,D,N}}) where {V,D,N} = TrackedStyle{N}() Broadcast.BroadcastStyle(::Type{<:TrackedReal}) = TrackedStyle{0}() -Broadcast.BroadcastStyle(::Type{<:AbstractArray{<:TrackedReal,N}}) where {N} = TrackedStyle{N}() +Broadcast.BroadcastStyle(::Type{<:AbstractArray{<:TrackedReal{<:Any,D},N}}) where {D,N} = + TrackedStyle{N}() # `AbstractArrayStyle{Any}` carries `Any` as its dimension, which `max` cannot compare _maxdim(M::Int, N::Int) = max(M, N) @@ -105,7 +105,6 @@ Base.copyto!(::TrackedArray, ::Broadcasted{<:TrackedStyle}) = _no_tracked_dest() Base.copyto!(::TrackedArray, ::Broadcasted{<:DefaultArrayStyle}) = _no_tracked_dest() getouttype(::TrackedReal{<:Any, D}) where {D} = D -getouttype(::TrackedArray{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D getouttype(::Any) = Union{} diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index f32396b2..9c6d8f13 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -199,6 +199,17 @@ end @test map(deriv, x) ≈ b end +@testset "an array of `TrackedReal`s with an abstract element type" begin + a = [1.0, 2.0] + + @test ReverseDiff.gradient(u -> sum(TrackedReal[u[1], u[2]] .* 2.0), a) == [2.0, 2.0] + @test ReverseDiff.gradient(u -> sum(exp.(TrackedReal[u[1], u[2]])), a) ≈ exp.(a) + @test ReverseDiff.gradient(u -> sum(u .* TrackedReal[u[2], u[1]]), a) == [4.0, 2.0] + + tape = ReverseDiff.GradientTape(u -> sum(exp.(TrackedReal[u[1], u[2]])), a) + @test ReverseDiff.gradient!(tape, [0.5, 1.5]) ≈ exp.([0.5, 1.5]) +end + @testset "functions constant in their argument" begin a = rand(3) From ce84f7b3c37e0bb14d9881ac9354d99e4ed95327 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:34:08 +0200 Subject: [PATCH 08/36] Read a known partial's arguments once per reverse pass `value` of an array of `TrackedReal`s or of an adjoint `TrackedVector` maps over it, so calling it per element made the reverse pass quadratic. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 3 ++- src/derivatives/propagation.jl | 2 +- test/derivatives/BroadcastTests.jl | 3 +++ 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 477f6286..0d7599c3 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -324,8 +324,9 @@ end results, _, tag, bounds = instruction.cache T = typeof(tag) slots, _ = trackedslots(input) + vals = map(value, input) map((x, slot, bound) -> - _br_add_to_deriv!(T, x, slot, output_deriv, results, bound, input), + _br_add_to_deriv!(T, x, slot, output_deriv, results, bound, vals), input, slots, bounds) unseed!(output) return nothing diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index 052a24ce..c8190bdf 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -108,7 +108,7 @@ struct Contract{Op,A<:Tuple} args::A end -_at(::Val{j}, i, args) where {j} = _elem(value(args[j]), i) +_at(::Val{j}, i, args) where {j} = _elem(args[j], i) _elem(v::Real, i) = v _elem(v::AbstractArray, i) = v[i] diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 9c6d8f13..806dff74 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -379,6 +379,9 @@ end @test ReverseDiff.gradient(t -> sum(t ./ w), v) == 1 ./ w @test ReverseDiff.gradient(t -> sum(w .\ t), v) == 1 ./ w @test ReverseDiff.gradient(t -> sum(t .* t), v) == 2 .* v + @test ReverseDiff.gradient(t -> sum(t .* Real[w...]), v) == w + @test ReverseDiff.gradient(t -> sum(t .* [t[i] for i in 1:3]), v) == 2 .* v + @test ReverseDiff.gradient(t -> sum(t' .* t'), v) == 2 .* v # `-x/y^2` is no argument of the broadcast, so it is read off a `Dual` @test ReverseDiff.gradient(t -> sum(w ./ t), v) == -w ./ v .^ 2 From 256c22fe6c034882d0085a251c8ed7c981233592 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:36:35 +0200 Subject: [PATCH 09/36] Drop a redundant `trackresults` method and two `_maxdim` methods Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 0d7599c3..a2b22e1e 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -52,9 +52,7 @@ Broadcast.BroadcastStyle(::Type{<:AbstractArray{<:TrackedReal{<:Any,D},N}}) wher # `AbstractArrayStyle{Any}` carries `Any` as its dimension, which `max` cannot compare _maxdim(M::Int, N::Int) = max(M, N) -_maxdim(::Type{Any}, ::Int) = Any -_maxdim(::Int, ::Type{Any}) = Any -_maxdim(::Type{Any}, ::Type{Any}) = Any +_maxdim(_, _) = Any # tracked values must stay tracked, so take precedence over every other array style Broadcast.BroadcastStyle(::TrackedStyle{M}, ::AbstractArrayStyle{N}) where {M,N} = @@ -284,9 +282,6 @@ end return out end -@inline trackresults(::Type{T}, results::AbstractArray{<:Dual{T}}, df, vf, targs, - ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, D) - @inline trackresults(::Type{T}, results::KnownPartials, df, vf, targs, ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, D) From a0708a895e575d2e40742bca6a6a92741c0a4ac7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:44:21 +0200 Subject: [PATCH 10/36] Name the closed-form reverse pass `contract_increment_deriv!` `diffresult_increment_deriv!` took either a partial's index or the arguments a `Contract` names in the same position. The replay and the tag check reuse the values `copy` already computed. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 48 ++++++++++++++---------------- src/derivatives/propagation.jl | 23 +++++++------- test/derivatives/BroadcastTests.jl | 6 ++-- 3 files changed, 39 insertions(+), 38 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index a2b22e1e..39ce0766 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -247,15 +247,14 @@ broadcastresults(entries::Tuple, slots, df, vals) = # known partials leave nothing to read off a `Dual`, so `f` is evaluated undualized vf(x::Vararg{Any,N}) where {N} = splatcall(f, x, untracked, inds) entries = knownpartials(f, args...) - return trackresults(T, broadcastresults(entries, slots, df, vals), df, vf, targs, D) + return trackresults(T, broadcastresults(entries, slots, df, vals), df, vf, targs, vals, D) end # the cache carries what the replay needs: `df` to recompute the stored partials, or, where # they are known already, `vf` for the values alone -replaycache(::Type{T}, results::AbstractArray, df, vf, targs) where {T} = +replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} = (df, map(y -> ForwardDiff.value(T, y), results)) -replaycache(::Type, ::KnownPartials, df, vf, targs) = - (vf, broadcast(vf, map(value, targs)...)) +replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) # the seed is contracted with an argument, so a perturbation riding on one ends up in the # derivative and `D` has to be able to hold it; a widened element type is a `Union` @@ -272,9 +271,10 @@ replaycache(::Type, ::KnownPartials, df, vf, targs) = end # `df` hands back `f`'s own result, so `results` keeps the element type `Base` would produce -@inline function recordresults(::Type{T}, results, df, vf, targs, ::Type{D}) where {T, D} - foreach(t -> checkargtags(D, eltype(value(t))), targs) - g, outvalue = replaycache(T, results, df, vf, targs) +@inline function recordresults(::Type{T}, results, df, vf, targs, vals, + ::Type{D}) where {T, D} + foreach(v -> checkargtags(D, eltype(v)), vals) + g, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) out = track(outvalue, D, tp) cache = (results, g, T(), map(t -> index_bound(t, out), targs)) @@ -282,8 +282,8 @@ end return out end -@inline trackresults(::Type{T}, results::KnownPartials, df, vf, targs, - ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, D) +@inline trackresults(::Type{T}, results::KnownPartials, df, vf, targs, vals, + ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, vals, D) # an enclosing differentiation's tag is constant in our arguments, while one nested inside # `f` buries our partial; a widened element type is a `Union`, so every member is checked @@ -302,13 +302,13 @@ end # a type-unstable `f`, such as one with an integer literal branch, leaves an abstract # element type that can still hide a `Dual{T}` -@inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, +@inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, vals, ::Type{D}) where {T, D} if typeintersect(eltype(results), Dual{T}) === Union{} checktags(T, eltype(results)) return results else - return recordresults(T, results, df, vf, targs, D) + return recordresults(T, results, df, vf, targs, vals, D) end end @@ -327,29 +327,27 @@ end return nothing end -# a per-element cache is indexed by the slot, a closed-form one is picked out by it -selectpartial(results, ::Val{k}, args) where {k} = (results, k) -selectpartial(p::KnownPartials, ::Val{k}, args) where {k} = (p.entries[k], args) - _br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::CartesianIndex, _) = nothing _br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::Nothing, _) = nothing # an argument broadcast to the full output shape needs no index clamping -function _br_add_to_deriv!(::Type{T}, x, slot::Val{k}, out_deriv, results, - bound::CartesianIndex, args) where {T, k} - results, sel = selectpartial(results, slot, args) +function _br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, + bound::CartesianIndex, args) where {T} if bound == CartesianIndex(size(out_deriv)) - return diffresult_increment_deriv!(T, x, out_deriv, results, sel) + return _increment_deriv!(T, x, out_deriv, results, slot, args) else - return diffresult_increment_deriv!(T, x, out_deriv, results, sel, bound) + return _increment_deriv!(T, x, out_deriv, results, slot, args, bound) end end -function _br_add_to_deriv!(::Type{T}, x, slot::Val{k}, out_deriv, results, ::Nothing, - args) where {T, k} - results, sel = selectpartial(results, slot, args) - return diffresult_increment_deriv!(T, x, out_deriv, results, sel, nothing) -end +_br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, ::Nothing, args) where {T} = + _increment_deriv!(T, x, out_deriv, results, slot, args, nothing) + +# a per-element cache is indexed by the slot, a closed-form one is picked out by it +_increment_deriv!(::Type{T}, x, out_deriv, results, ::Val{k}, args, bound...) where {T, k} = + diffresult_increment_deriv!(T, x, out_deriv, results, k, bound...) +_increment_deriv!(::Type, x, out_deriv, p::KnownPartials, ::Val{k}, args, bound...) where {k} = + contract_increment_deriv!(x, out_deriv, p.entries[k], args, bound...) @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) input, output = instruction.input, instruction.output diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index c8190bdf..c4aa5ef6 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -143,28 +143,31 @@ function diffresult_increment_deriv!(::Type{T}, input::TrackedReal, x::AbstractA return nothing end -# a closed-form partial is named rather than stored, so `args` stands where an index would -function diffresult_increment_deriv!(::Type, input::AbstractArray, x::AbstractArray, - results::Contract, args::Tuple) +############################# +# contract_increment_deriv! # +############################# + +function contract_increment_deriv!(input::AbstractArray, x::AbstractArray, e::Contract, + args::Tuple) for i in eachindex(input, x) - increment_deriv!(input, _contract(results, x[i], i, args), i) + increment_deriv!(input, _contract(e, x[i], i, args), i) end return nothing end -function diffresult_increment_deriv!(::Type, input::AbstractArray, x::AbstractArray, - results::Contract, args::Tuple, bound::CartesianIndex) +function contract_increment_deriv!(input::AbstractArray, x::AbstractArray, e::Contract, + args::Tuple, bound::CartesianIndex) for i in CartesianIndices(size(x)) - increment_deriv!(input, _contract(results, x[i], i, args), min(bound, i)) + increment_deriv!(input, _contract(e, x[i], i, args), min(bound, i)) end return nothing end -function diffresult_increment_deriv!(::Type, input::TrackedReal, x::AbstractArray, - results::Contract, args::Tuple, ::Nothing) +function contract_increment_deriv!(input::TrackedReal, x::AbstractArray, e::Contract, + args::Tuple, ::Nothing) isempty(x) && return nothing pull_deriv!(input) - input.deriv += sum(i -> _contract(results, x[i], i, args), eachindex(x)) + input.deriv += sum(i -> _contract(e, x[i], i, args), eachindex(x)) push_deriv!(input) return nothing end diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 806dff74..a5da1f46 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -257,10 +257,10 @@ end 1.0, ForwardDiff.Dual{tagA}(2.0, 3.0)] nested = Union{Float64,ForwardDiff.Dual{tagB,Float64,1}}[1.0] - @test ReverseDiff.trackresults(tagB, concrete, identity, identity, (), Float64) === concrete - @test ReverseDiff.trackresults(tagB, widened, identity, identity, (), Float64) === widened + @test ReverseDiff.trackresults(tagB, concrete, identity, identity, (), (), Float64) === concrete + @test ReverseDiff.trackresults(tagB, widened, identity, identity, (), (), Float64) === widened @test_throws ForwardDiff.DualMismatchError ReverseDiff.trackresults( - tagA, nested, identity, identity, (), Float64) + tagA, nested, identity, identity, (), (), Float64) @test ReverseDiff.getpartial(tagA, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 2.0 @test ReverseDiff.getpartial(tagA, 1.0, 1) == 0.0 From b7fdb918800d9d65b5b097997dbe728b4b1af4e0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Wed, 23 Sep 2026 22:45:47 +0200 Subject: [PATCH 11/36] Link the perturbation-confusion TODO to #45 Co-Authored-By: Claude Opus 5.5 (1M context) --- src/tracked.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tracked.jl b/src/tracked.jl index e0486947..fc687e8a 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -518,7 +518,7 @@ track(x::Real, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = T track(x::AbstractArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = TrackedArray(x, fill!(similar(x, D), zero(D)), tp) # every forward pass writes into the value buffer, which a `TrackedArray` rejects -# TODO: self-nesting, and the `value`/`track` unwrapping elsewhere, risk perturbation confusion +# TODO: self-nesting, and the `value`/`track` unwrapping elsewhere, risk perturbation confusion (#45) track(x::TrackedArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = track(collect(x), D, tp) track!(t::TrackedArray, x::AbstractArray) = (value!(t, x); unseed!(t); t) From d37b386d599feeb5d3d85f229fc89ddff3de4091 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 24 Sep 2026 08:59:52 +0200 Subject: [PATCH 12/36] Handle uninstantiated broadcasts in `copy` (fixes tests on Julia 1.13) LinearAlgebra 1.13 forwards a `Diagonal`'s broadcast to its diagonal and calls `copy` on the result without instantiating it. With a tracked diagonal, the `Broadcasted` has `TrackedStyle{1}` and no axes, which the scalar `copy` method matched by its axes type, so it was indexed with `CartesianIndex()` and threw a `BoundsError`. Decide the scalar route from `axes` instead, which computes missing axes and still gives `()` for `TrackedStyle{0}` and 0-dimensional `TrackedStyle{Any}`. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 11 ++++++----- test/derivatives/BroadcastTests.jl | 19 +++++++++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 39ce0766..3df8bdfc 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -77,11 +77,12 @@ function remove_not_tracked(b::Broadcasted{style}) where {style} return Broadcasted{style}(remove_not_tracked(b.f), remove_not_tracked.(b.args), b.axes) end -# scalars take `Base`'s 0-dimensional route onto the scalar derivative rules; `instantiate` -# leaves an `AbstractArrayStyle{0}` without axes, so 0 dimensions shows up either way -Base.copy(bc::Broadcasted{<:TrackedStyle, <:Union{Nothing, Tuple{}}}) = bc[CartesianIndex()] - function Base.copy(_bc::Broadcasted{<:TrackedStyle}) + # scalars take the scalar derivative rules; ask the axes, since `TrackedStyle{Any}` carries + # no dimension and `LinearAlgebra` may pass an uninstantiated `Broadcasted` + if axes(_bc) isa Tuple{} + return _bc[CartesianIndex()] + end bc = remove_not_tracked(_bc) flattened_bc = Base.Broadcast.flatten(bc) f, args = flattened_bc.f, flattened_bc.args @@ -231,7 +232,7 @@ broadcastresults(entries::Tuple, slots, df, vals) = KnownPartials(trackedentries(slots, entries)) # at least one argument has to be a non-0-dimensional array: `copy` sends the scalar and -# 0-dimensional cases down `Base`'s `TrackedStyle{0}` route onto the scalar rules instead +# 0-dimensional cases onto the scalar rules instead @inline function ∇broadcast(f::F, args::Tuple, argvals::Tuple) where {F} inds, targs, untracked = splitargs(args) _, vals, _ = splitargs(argvals) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index a5da1f46..27e5d362 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -325,6 +325,25 @@ end @test ReverseDiff.gradient(v -> sum(v .* Foreign0d(2.0)), [1.0, 2.0]) ≈ [2.0, 2.0] end +@testset "a `Broadcasted` that reaches `copy` uninstantiated" begin + # `LinearAlgebra` forwards a `Diagonal`'s broadcast to its diagonal without instantiating it + a = rand(3) + tp = InstructionTape() + x, tr = track(copy(a), tp), track(2.0, tp) + + @test value(copy(Broadcast.broadcasted(exp, x))) ≈ exp.(a) + @test copy(Broadcast.broadcasted(exp, tr)) isa TrackedReal + @test copy(Broadcast.broadcasted(+, tr, Foreign0d(1.0))) isa TrackedReal + + @test ReverseDiff.gradient(v -> sum(Diagonal(v) .* 2.0), a) == fill(2.0, 3) + + # alongside a `TrackedArray`, the wrapper is an argument of `∇broadcast` + m = rand(3, 3) + gv, gm = ReverseDiff.gradient((v, w) -> sum(Diagonal(v) .* w), (a, m)) + @test gv == diag(m) + @test gm == Diagonal(a) +end + @testset "`NotTracked`" begin f = ReverseDiff.NotTracked(t -> 2t) From 6b23a2843b62c87f6ab84b46db8f50b1d131a7c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 24 Sep 2026 14:41:50 +0200 Subject: [PATCH 13/36] =?UTF-8?q?Strip=20`=E2=88=87broadcast`'s=20`Dual`s?= =?UTF-8?q?=20in=20`value`=20(fixes=20tests=20on=20Julia=201.13)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `∇broadcast` now seeds its `Dual`s with a ReverseDiff tag, and `value` strips them as it strips tracking. So a broadcast of `value`, fused or not, isn't recorded. On Julia 1.13 `map(value, ::Diagonal)` broadcasts over the diagonal, which used to keep `value` of a tracked `Diagonal` tracked. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/broadcast.jl | 8 +++++++- test/derivatives/BroadcastTests.jl | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 3df8bdfc..aaaab9c8 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -231,6 +231,12 @@ broadcastresults(::Nothing, slots, df, vals) = broadcast(df, vals...) broadcastresults(entries::Tuple, slots, df, vals) = KnownPartials(trackedentries(slots, entries)) +# marks the `Dual`s seeded by `∇broadcast`, so `value` can strip them as it strips tracking; +# `F` gives each function its own tag, which `ForwardDiff` orders by first use, as `Tag(f, V)` would +struct BroadcastTag{F} end + +value(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = ForwardDiff.value(T, x) + # at least one argument has to be a non-0-dimensional array: `copy` sends the scalar and # 0-dimensional cases onto the scalar rules instead @inline function ∇broadcast(f::F, args::Tuple, argvals::Tuple) where {F} @@ -239,7 +245,7 @@ broadcastresults(entries::Tuple, slots, df, vals) = D = mapreduce(getouttype, promote_type, targs) slots, valP = trackedslots(targs) # one tag for the whole broadcast keeps `results` concretely typed - T = typeof(ForwardDiff.Tag(f, reduce(promote_type, map(dualvaltype, slots, vals)))) + T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), reduce(promote_type, map(dualvaltype, slots, vals)))) # `broadcast` calls `df` elementwise, so it receives one scalar per argument function df(x::Vararg{Any,N}) where {N} dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 27e5d362..632c3312 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -344,6 +344,23 @@ end @test gm == Diagonal(a) end +@testset "`value` in a broadcast is not recorded" begin + a = rand(3) + tp = InstructionTape() + x = track(copy(a), tp) + + d = value(Diagonal(x)) + @test d isa Diagonal{Float64,Vector{Float64}} + @test d == Diagonal(a) + @test isempty(tp) + + # also when fused with a differentiable function + y = value.(x .+ 1) + @test y isa Vector{Float64} + @test y ≈ a .+ 1 + @test isempty(tp) +end + @testset "`NotTracked`" begin f = ReverseDiff.NotTracked(t -> 2t) From ae10d2a5f4ec4f94d9ced45944cbd628c373b446 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 24 Sep 2026 21:28:15 +0200 Subject: [PATCH 14/36] Expect `@skip` broadcasts to be untracked, as on master MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since `value` strips `∇broadcast`'s `Dual`s, `SkipOptimize` does too, so a broadcast of a skipped function is neither differentiated nor recorded. This is what master does with `broadcast(f::SkipOptimize, t::TrackedArray) = broadcast(f.f, value(t))`, and what `map` and the scalar rules do. Co-Authored-By: Claude Opus 5.5 (1M context) --- test/derivatives/BroadcastTests.jl | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 632c3312..9406248d 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -111,15 +111,19 @@ end end @testset "`@forward` and `@skip` wrappers do not force the fallback path" begin - for wrapper in (ReverseDiff.ForwardOptimize, ReverseDiff.SkipOptimize) - tp = InstructionTape() - x = track(rand(3, 3), tp) - - y = broadcast(wrapper(exp), x) + tp = InstructionTape() + x = track(rand(3, 3), tp) + y = broadcast(ReverseDiff.ForwardOptimize(exp), x) + @test y isa TrackedArray + @test length(tp) == 1 - @test y isa TrackedArray - @test length(tp) == 1 - end + # `@skip` results are untracked, as with `map` and scalars + tp = InstructionTape() + x = track(rand(3, 3), tp) + y = broadcast(ReverseDiff.SkipOptimize(exp), x) + @test y isa Matrix{Float64} + @test y == exp.(value(x)) + @test isempty(tp) end @testset "a type broadcast as a function does not force the fallback path" begin From f370cb2b9494b0d82889d72e61a37a2343ec50a3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 25 Sep 2026 14:28:40 +0200 Subject: [PATCH 15/36] Read argument values in the reverse pass only for known partials `value` copies an array of `TrackedReal`s, so the `Dual` path allocated a copy per argument on every reverse pass without reading it. Also make two tests exercise what they describe: a `Ref` that shifts argument positions, and a replay whose named argument is tracked. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 35 +++++++++++++++++++----------- test/derivatives/BroadcastTests.jl | 9 ++++---- 2 files changed, 27 insertions(+), 17 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index aaaab9c8..d8d3237e 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -326,35 +326,44 @@ end results, _, tag, bounds = instruction.cache T = typeof(tag) slots, _ = trackedslots(input) - vals = map(value, input) - map((x, slot, bound) -> - _br_add_to_deriv!(T, x, slot, output_deriv, results, bound, vals), + partials = reversepartials(results, input) + map((x, slot, bound) -> _br_add_to_deriv!(T, x, slot, output_deriv, partials, bound), input, slots, bounds) unseed!(output) return nothing end -_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::CartesianIndex, _) = nothing -_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::Nothing, _) = nothing +struct PartialsWithArgs{E<:Tuple,A<:Tuple} + entries::E + args::A +end + +# the partials of one reverse pass. Only known partials read the argument values, so the +# per-element cache skips computing them, which copies an array of `TrackedReal`s. +reversepartials(results::AbstractArray, _) = results +reversepartials(p::KnownPartials, input) = PartialsWithArgs(p.entries, map(value, input)) + +_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::CartesianIndex) = nothing +_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::Nothing) = nothing # an argument broadcast to the full output shape needs no index clamping function _br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, - bound::CartesianIndex, args) where {T} + bound::CartesianIndex) where {T} if bound == CartesianIndex(size(out_deriv)) - return _increment_deriv!(T, x, out_deriv, results, slot, args) + return _increment_deriv!(T, x, out_deriv, results, slot) else - return _increment_deriv!(T, x, out_deriv, results, slot, args, bound) + return _increment_deriv!(T, x, out_deriv, results, slot, bound) end end -_br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, ::Nothing, args) where {T} = - _increment_deriv!(T, x, out_deriv, results, slot, args, nothing) +_br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, ::Nothing) where {T} = + _increment_deriv!(T, x, out_deriv, results, slot, nothing) # a per-element cache is indexed by the slot, a closed-form one is picked out by it -_increment_deriv!(::Type{T}, x, out_deriv, results, ::Val{k}, args, bound...) where {T, k} = +_increment_deriv!(::Type{T}, x, out_deriv, results, ::Val{k}, bound...) where {T, k} = diffresult_increment_deriv!(T, x, out_deriv, results, k, bound...) -_increment_deriv!(::Type, x, out_deriv, p::KnownPartials, ::Val{k}, args, bound...) where {k} = - contract_increment_deriv!(x, out_deriv, p.entries[k], args, bound...) +_increment_deriv!(::Type, x, out_deriv, p::PartialsWithArgs, ::Val{k}, bound...) where {k} = + contract_increment_deriv!(x, out_deriv, p.entries[k], p.args, bound...) @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) input, output = instruction.input, instruction.output diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 9406248d..f14654b2 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -436,7 +436,8 @@ end @test ReverseDiff.gradient(t -> sum(t .* t[1]), v) == [2 * v[1] + sum(v[2:end]), 1.0, 1.0] # an argument `splitargs` holds back is not one a partial may name - @test ReverseDiff.gradient(t -> sum(broadcast(*, t, w, Ref(2.0))), v) == 2 .* w + @test ReverseDiff.gradient(t -> sum(Ref(2.0) .* t), v) == fill(2.0, 3) + @test ReverseDiff.gradient(t -> sum(t .* Ref(2.0)), v) == fill(2.0, 3) ga, gb = ReverseDiff.gradient((a, b) -> sum(a .- b), (v, w)) @test ga == fill(1.0, 3) @@ -458,9 +459,9 @@ end @test ReverseDiff.gradient!(tape, x) ≈ 2 .* exp.(2 .* x) # an argument a partial names is read at replay time, not at record time - scaled(t) = t .* [2.0, 3.0, 4.0] - tape2 = ReverseDiff.GradientTape(t -> sum(scaled(t)), [1.0, 2.0, 3.0]) - @test ReverseDiff.gradient!(tape2, x) == [2.0, 3.0, 4.0] + tape2 = ReverseDiff.GradientTape(t -> sum(t .* t), [1.0, 2.0, 3.0]) + @test ReverseDiff.gradient!(tape2, x) == 2 .* x + @test ReverseDiff.gradient!(ReverseDiff.compile(tape2), x) == 2 .* x end end # module From cb5be0f38423054afc43d2568a395631d21cd057 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 25 Sep 2026 14:28:40 +0200 Subject: [PATCH 16/36] Include test files by relative path Co-Authored-By: Claude Opus 5.5 --- test/runtests.jl | 30 ++++++++++++++---------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index f5d4bde4..c5d7c174 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,56 +1,54 @@ using Test -const TESTDIR = dirname(@__FILE__) - test_println(kind, f, pad=" ") = println(pad, "testing $(kind): `$(f)`...") @testset "ReverseDiff" begin @testset "TapeTests" begin println("running TapeTests...") - t = @elapsed include(joinpath(TESTDIR, "TapeTests.jl")) + t = @elapsed include("TapeTests.jl") println("done (took $t seconds).") end @testset "TrackedTests" begin println("running TrackedTests...") - t = @elapsed include(joinpath(TESTDIR, "TrackedTests.jl")) + t = @elapsed include("TrackedTests.jl") println("done (took $t seconds).") end @testset "MacrosTests" begin println("running MacrosTests...") - t = @elapsed include(joinpath(TESTDIR, "MacrosTests.jl")) + t = @elapsed include("MacrosTests.jl") println("done (took $t seconds).") end @testset "ChainRulesTests" begin println("running ChainRulesTests...") - t = @elapsed include(joinpath(TESTDIR, "ChainRulesTests.jl")) + t = @elapsed include("ChainRulesTests.jl") println("done (took $t seconds).") end @testset "ScalarTests" begin println("running ScalarTests...") - t = @elapsed include(joinpath(TESTDIR, "derivatives/ScalarTests.jl")) + t = @elapsed include("derivatives/ScalarTests.jl") println("done (took $t seconds).") end @testset "LinAlgTests" begin println("running LinAlgTests...") - t = @elapsed include(joinpath(TESTDIR, "derivatives/LinAlgTests.jl")) + t = @elapsed include("derivatives/LinAlgTests.jl") println("done (took $t seconds).") end @testset "PrecisionTests" begin println("running PrecisionTests...") - t = @elapsed include(joinpath(TESTDIR, "PrecisionTests.jl")) + t = @elapsed include("PrecisionTests.jl") println("done (took $t seconds).") end @testset "ElementWiseTests" begin println("running ElementwiseTests...") - t = @elapsed include(joinpath(TESTDIR, "derivatives/ElementwiseTests.jl")) + t = @elapsed include("derivatives/ElementwiseTests.jl") println("done (took $t seconds).") end @@ -62,37 +60,37 @@ test_println(kind, f, pad=" ") = println(pad, "testing $(kind): `$(f)`...") @testset "ArrayFunctionTests" begin println("running ArrayFunctionTests...") - t = @elapsed include(joinpath(TESTDIR, "derivatives/ArrayFunctionTests.jl")) + t = @elapsed include("derivatives/ArrayFunctionTests.jl") println("done (took $t seconds).") end @testset "GradientTests" begin println("running GradientTests...") - t = @elapsed include(joinpath(TESTDIR, "api/GradientTests.jl")) + t = @elapsed include("api/GradientTests.jl") println("done (took $t seconds).") end @testset "JacobianTests" begin println("running JacobianTests...") - t = @elapsed include(joinpath(TESTDIR, "api/JacobianTests.jl")) + t = @elapsed include("api/JacobianTests.jl") println("done (took $t seconds).") end @testset "HessianTests" begin println("running HessianTests...") - t = @elapsed include(joinpath(TESTDIR, "api/HessianTests.jl")) + t = @elapsed include("api/HessianTests.jl") println("done (took $t seconds).") end @testset "ConfigTests" begin println("running ConfigTests...") - t = @elapsed include(joinpath(TESTDIR, "api/ConfigTests.jl")) + t = @elapsed include("api/ConfigTests.jl") println("done (took $t seconds).") end @testset "CompatTests" begin println("running CompatTests...") - t = @elapsed include(joinpath(TESTDIR, "compat/CompatTests.jl")) + t = @elapsed include("compat/CompatTests.jl") println("done (took $t seconds).") end end From a9a2b37ea858137c29a73e56c9e8f1f5c377d1d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 25 Sep 2026 15:33:50 +0200 Subject: [PATCH 17/36] Index broadcast partials by tracked argument only `trackedslots` also returns the position of each tracked argument, so known partials, bounds and the reverse pass cover tracked arguments alone instead of carrying `nothing` or `Val(0)` for untracked ones. Known partials are given per argument, and the tag is keyed on the derivative type. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 131 ++++++++++++++++------------------- 1 file changed, 61 insertions(+), 70 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index d8d3237e..0a1b1cda 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -143,18 +143,22 @@ end @inline incr(::Val{k}) where {k} = Val(k + 1) -# `slots[i]` indexes argument `i`'s partial and is `Val(0)` where the argument is untracked; -# the second value counts the tracked arguments. Both are built on the way out of the -# recursion: a count passed into it would stop being constant-folded after a few arguments, -# leaving the partials a runtime-sized tuple. -@inline trackedslots(::Tuple{}) = (), Val(0) +# argument `i` has partial `slots[i]` (`Val(0)` if untracked), partial `k` belongs to argument +# `positions[k]`. Built on the way out of the recursion to stay constant-folded. +@inline trackedslots(::Tuple{}) = (), (), Val(0) @inline function trackedslots(args::Tuple) - slots, n = trackedslots(Base.tail(args)) - istracked(first(args)) || return (Val(0), slots...), n - k = incr(n) - return (k, slots...), k + slots, positions, n = trackedslots(Base.tail(args)) + positions = map(incr, positions) + if istracked(first(args)) + k = incr(n) + return (k, slots...), (positions..., Val(1)), k + else + return (Val(0), slots...), positions, n + end end +@inline getat(xs::Tuple, ::Val{i}) where {i} = xs[i] + # `DiffRules` leaves partials such as the order of `besselj` as `NaN`, which `NaN * 0` # would spread to every slot, so untracked arguments are not dualized @inline dualize(::Type, ::Val{0}, ::Val, x) = x @@ -162,75 +166,59 @@ end return Dual{T}(x, ntuple(j -> j == k, valP)) end -# only seeded arguments make up the `Dual` value type, so the tag follows `slots` -@inline dualvaltype(::Val{0}, v) = Union{} -@inline dualvaltype(::Val, v) = eltype(v) - -# `f`'s partials in closed form, one entry per tracked argument in slot order +# `f`'s partials in closed form, one per tracked argument struct KnownPartials{E<:Tuple} entries::E end -# `trackedslots` numbers the tracked arguments from the last, so drop the untracked ones and -# reverse what remains -trackedentries(::Tuple{}, ::Tuple{}) = () -trackedentries(slots::Tuple{Val{0},Vararg{Any}}, entries::Tuple) = - trackedentries(Base.tail(slots), Base.tail(entries)) -trackedentries(slots::Tuple, entries::Tuple) = - (trackedentries(Base.tail(slots), Base.tail(entries))..., first(entries)) +const RealOrArray = Union{Real, AbstractArray{<:Real}} # an argument read at the output's own index has to span the whole output _sameshape(x, y) = x isa Real || y isa Real || axes(x) == axes(y) -# one entry per argument, or `nothing` where `f`'s partials depend on the point. Annotating -# every argument is what lets an entry name one: those are exactly the ones `splitargs` keeps. -knownpartials(f, args...) = nothing - -# `+` and `-` are affine in their arguments jointly -knownpartials(::Union{typeof(+), typeof(identity)}, - args::Vararg{Union{Real, AbstractArray{<:Real}}}) = - map(_ -> Contract(identity, ()), args) - -knownpartials(::typeof(-), arg::Union{Real, AbstractArray{<:Real}}) = (Contract(-, ()),) - -knownpartials(::typeof(-), x::Union{Real, AbstractArray{<:Real}}, - y::Union{Real, AbstractArray{<:Real}}) = - (Contract(identity, ()), Contract(-, ())) - -# `*`, `/` and `\` are multilinear, so a partial is another argument, which the reverse pass -# reads from the instruction's own input where `record!` has captured it. A tracked -# denominator is excluded, its partial `-x/y^2` being no argument of the broadcast. -function knownpartials(::typeof(*), x::Union{Real, AbstractArray{<:Real}}, - y::Union{Real, AbstractArray{<:Real}}) +function ifsameshape(c::Contract, x, y) if _sameshape(x, y) - return (Contract(*, (Val(2),)), Contract(*, (Val(1),))) + return c else return nothing end end -function knownpartials(::typeof(/), x::Union{Real, AbstractArray{<:Real}}, - y::Union{Real, AbstractArray{<:Real}}) - if !istracked(y) && _sameshape(x, y) - return (Contract(/, (Val(2),)), nothing) - else +# partial with respect to argument `i`, or `nothing`. Annotating every argument ensures that +# a named argument is one `splitargs` keeps. +knownpartial(f, ::Val, args) = nothing + +knownpartial(::Union{typeof(+), typeof(identity)}, ::Val, ::Tuple{Vararg{RealOrArray}}) = + Contract(identity, ()) +knownpartial(::typeof(-), ::Val{1}, ::Tuple{RealOrArray}) = Contract(-, ()) +knownpartial(::typeof(-), ::Val{1}, ::Tuple{RealOrArray,RealOrArray}) = Contract(identity, ()) +knownpartial(::typeof(-), ::Val{2}, ::Tuple{RealOrArray,RealOrArray}) = Contract(-, ()) + +# a denominator's partial `-x/y^2` is no argument of the broadcast +knownpartial(::typeof(*), ::Val{1}, (x, y)::Tuple{RealOrArray,RealOrArray}) = + ifsameshape(Contract(*, (Val(2),)), x, y) +knownpartial(::typeof(*), ::Val{2}, (x, y)::Tuple{RealOrArray,RealOrArray}) = + ifsameshape(Contract(*, (Val(1),)), x, y) +knownpartial(::typeof(/), ::Val{1}, (x, y)::Tuple{RealOrArray,RealOrArray}) = + ifsameshape(Contract(/, (Val(2),)), x, y) +knownpartial(::typeof(\), ::Val{2}, (x, y)::Tuple{RealOrArray,RealOrArray}) = + ifsameshape(Contract(/, (Val(1),)), x, y) + +# `nothing` unless every tracked argument's partial is known +knownpartials(f, args, ::Tuple{}) = () +function knownpartials(f, args, positions::Tuple) + rest = knownpartials(f, args, Base.tail(positions)) + if rest === nothing return nothing end -end - -function knownpartials(::typeof(\), x::Union{Real, AbstractArray{<:Real}}, - y::Union{Real, AbstractArray{<:Real}}) - if !istracked(x) && _sameshape(x, y) - return (nothing, Contract(/, (Val(1),))) - else + p = knownpartial(f, first(positions), args) + if p === nothing return nothing + else + return (p, rest...) end end -broadcastresults(::Nothing, slots, df, vals) = broadcast(df, vals...) -broadcastresults(entries::Tuple, slots, df, vals) = - KnownPartials(trackedentries(slots, entries)) - # marks the `Dual`s seeded by `∇broadcast`, so `value` can strip them as it strips tracking; # `F` gives each function its own tag, which `ForwardDiff` orders by first use, as `Tag(f, V)` would struct BroadcastTag{F} end @@ -243,9 +231,9 @@ value(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = ForwardDiff.value inds, targs, untracked = splitargs(args) _, vals, _ = splitargs(argvals) D = mapreduce(getouttype, promote_type, targs) - slots, valP = trackedslots(targs) - # one tag for the whole broadcast keeps `results` concretely typed - T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), reduce(promote_type, map(dualvaltype, slots, vals)))) + slots, positions, valP = trackedslots(targs) + # one tag per broadcast, and a new one under an enclosing differentiation + T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), D)) # `broadcast` calls `df` elementwise, so it receives one scalar per argument function df(x::Vararg{Any,N}) where {N} dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) @@ -253,8 +241,13 @@ value(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = ForwardDiff.value end # known partials leave nothing to read off a `Dual`, so `f` is evaluated undualized vf(x::Vararg{Any,N}) where {N} = splatcall(f, x, untracked, inds) - entries = knownpartials(f, args...) - return trackresults(T, broadcastresults(entries, slots, df, vals), df, vf, targs, vals, D) + entries = knownpartials(f, args, positions) + if entries === nothing + results = broadcast(df, vals...) + else + results = KnownPartials(entries) + end + return trackresults(T, results, df, vf, targs, vals, D) end # the cache carries what the replay needs: `df` to recompute the stored partials, or, where @@ -284,7 +277,8 @@ end g, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) out = track(outvalue, D, tp) - cache = (results, g, T(), map(t -> index_bound(t, out), targs)) + _, positions, _ = trackedslots(targs) + cache = (results, g, T(), map(p -> index_bound(getat(targs, p), out), positions)) record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache) return out end @@ -325,10 +319,10 @@ end output_deriv = deriv(output) results, _, tag, bounds = instruction.cache T = typeof(tag) - slots, _ = trackedslots(input) + _, positions, n = trackedslots(input) partials = reversepartials(results, input) - map((x, slot, bound) -> _br_add_to_deriv!(T, x, slot, output_deriv, partials, bound), - input, slots, bounds) + map((p, k, bound) -> _br_add_to_deriv!(T, getat(input, p), k, output_deriv, partials, bound), + positions, ntuple(Val, n), bounds) unseed!(output) return nothing end @@ -343,9 +337,6 @@ end reversepartials(results::AbstractArray, _) = results reversepartials(p::KnownPartials, input) = PartialsWithArgs(p.entries, map(value, input)) -_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::CartesianIndex) = nothing -_br_add_to_deriv!(::Type, _, ::Val{0}, _, _, ::Nothing) = nothing - # an argument broadcast to the full output shape needs no index clamping function _br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, bound::CartesianIndex) where {T} From 9217f196de1139fa43094d9c23aa22782bbc1726 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 25 Sep 2026 23:45:39 +0200 Subject: [PATCH 18/36] Keep tapeless tracked scalars constant and fuse constant tuples `capture` records a `TrackedReal` without a tape as its value, so the reverse pass recounting tracked inputs misaligned slots and bounds. The cache now carries each tracked argument's position, slot and bound from recording, and the reverse pass skips arguments captured untracked. Tuple arguments are broadcast elementwise rather than handed to `f` whole, and only fall back when an element is tracked or may hide one. Also unroll the vestigial one-element `map` loops, simplify comments, and make the elementwise test state which functions are left off the tape. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 47 +++---- src/derivatives/elementwise.jl | 185 +++++++++++++-------------- src/derivatives/propagation.jl | 2 +- test/derivatives/BroadcastTests.jl | 44 ++++++- test/derivatives/ElementwiseTests.jl | 6 +- 5 files changed, 154 insertions(+), 130 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 0a1b1cda..2ac5eac3 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -15,14 +15,14 @@ struct NotTracked{F} <: Function end (f::NotTracked{<:Union{Function, Type}})(args...; kwargs...) = f.f(args...; kwargs...) -# an argument is seeded itself, or elementwise, so its own tracked values are in plain sight; a -# `Ref` is handed to `f` whole, so nothing inside one is -mayhidetracked(b::F) where {F} = _mayhidetracked(F) +# can `f` receive tracked values that `∇broadcast` does not seed? +mayhidetracked(::F) where {F} = _mayhidetracked(F) mayhidetracked(::AbstractArray{F}) where {F} = _mayhidetracked(F) mayhidetracked(::Base.RefValue{F}) where {F} = _mayhidetracked(F) mayhidetracked(::AbstractArray{<:Real}) = false mayhidetracked(::Real) = false -mayhidetracked(b::Type) = false +mayhidetracked(::Type) = false +mayhidetracked(t::Tuple) = any(x -> istracked(x) || mayhidetracked(x), t) mayhidetracked(b::ForwardOptimize) = mayhidetracked(b.f) mayhidetracked(b::SkipOptimize) = mayhidetracked(b.f) mayhidetracked(b::Broadcasted) = mayhidetracked(b.f) || any(mayhidetracked, b.args) @@ -131,8 +131,7 @@ end @generated function splitargs(args::T) where {T <: Tuple} N = length(T.types) - RealOrArray = Union{Real, AbstractArray} - inds = [i for i in 1:N if T.types[i] <: RealOrArray] + inds = [i for i in 1:N if T.types[i] <: Union{Real, AbstractArray, Tuple}] indsval = :(Val{$(Expr(:tuple, [:($i) for i in inds]...))}()) maybetracked = Expr(:tuple, [:(args[$i]) for i in inds]...) untracked = Expr(:tuple, [:(args[$i]) for i in 1:N if !(i in inds)]...) @@ -144,7 +143,8 @@ end @inline incr(::Val{k}) where {k} = Val(k + 1) # argument `i` has partial `slots[i]` (`Val(0)` if untracked), partial `k` belongs to argument -# `positions[k]`. Built on the way out of the recursion to stay constant-folded. +# `positions[k]`, numbered from the last tracked argument. Built on the way out of the recursion +# to stay constant-folded. @inline trackedslots(::Tuple{}) = (), (), Val(0) @inline function trackedslots(args::Tuple) slots, positions, n = trackedslots(Base.tail(args)) @@ -174,18 +174,10 @@ end const RealOrArray = Union{Real, AbstractArray{<:Real}} # an argument read at the output's own index has to span the whole output -_sameshape(x, y) = x isa Real || y isa Real || axes(x) == axes(y) +ifsameshape(c::Contract, x, y) = x isa Real || y isa Real || axes(x) == axes(y) ? c : nothing -function ifsameshape(c::Contract, x, y) - if _sameshape(x, y) - return c - else - return nothing - end -end - -# partial with respect to argument `i`, or `nothing`. Annotating every argument ensures that -# a named argument is one `splitargs` keeps. +# partial with respect to argument `i`, or `nothing`. Requiring every argument to be +# `RealOrArray` makes positions in `args` positions in the arguments `splitargs` keeps. knownpartial(f, ::Val, args) = nothing knownpartial(::Union{typeof(+), typeof(identity)}, ::Val, ::Tuple{Vararg{RealOrArray}}) = @@ -256,8 +248,8 @@ replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} = (df, map(y -> ForwardDiff.value(T, y), results)) replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) -# the seed is contracted with an argument, so a perturbation riding on one ends up in the -# derivative and `D` has to be able to hold it; a widened element type is a `Union` +# a perturbation riding on an argument ends up in the derivative, so `D` has to be able to hold +# it; a widened element type is a `Union` @inline function checkargtags(::Type{D}, ::Type{E}) where {D, E} if E isa Union checkargtags(D, E.a) @@ -270,15 +262,15 @@ replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) return nothing end -# `df` hands back `f`'s own result, so `results` keeps the element type `Base` would produce @inline function recordresults(::Type{T}, results, df, vf, targs, vals, ::Type{D}) where {T, D} foreach(v -> checkargtags(D, eltype(v)), vals) g, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) out = track(outvalue, D, tp) - _, positions, _ = trackedslots(targs) - cache = (results, g, T(), map(p -> index_bound(getat(targs, p), out), positions)) + _, positions, n = trackedslots(targs) + bounds = map(p -> index_bound(getat(targs, p), out), positions) + cache = (results, g, T(), map(tuple, positions, ntuple(Val, n), bounds)) record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache) return out end @@ -317,12 +309,13 @@ end input = instruction.input output = instruction.output output_deriv = deriv(output) - results, _, tag, bounds = instruction.cache + results, _, tag, targets = instruction.cache T = typeof(tag) - _, positions, n = trackedslots(input) partials = reversepartials(results, input) - map((p, k, bound) -> _br_add_to_deriv!(T, getat(input, p), k, output_deriv, partials, bound), - positions, ntuple(Val, n), bounds) + map(targets) do (p, k, bound) + x = getat(input, p) + istracked(x) && _br_add_to_deriv!(T, x, k, output_deriv, partials, bound) + end unseed!(output) return nothing end diff --git a/src/derivatives/elementwise.jl b/src/derivatives/elementwise.jl index 8c3dfaf8..35c58f0a 100644 --- a/src/derivatives/elementwise.jl +++ b/src/derivatives/elementwise.jl @@ -5,28 +5,28 @@ # dispatch # #----------# -for g in (:map,), f in SKIPPED_UNARY_SCALAR_FUNCS - @eval @inline Base.$(g)(f::typeof($f), t::TrackedArray) = $(g)(SkipOptimize(f), t) +for f in SKIPPED_UNARY_SCALAR_FUNCS + @eval @inline Base.map(f::typeof($f), t::TrackedArray) = map(SkipOptimize(f), t) end -for g in (:map,), f in SKIPPED_BINARY_SCALAR_FUNCS +for f in SKIPPED_BINARY_SCALAR_FUNCS @eval begin - @inline Base.$(g)(f::typeof($f), x::TrackedArray, y::TrackedArray) = $(g)(SkipOptimize(f), x, y) - @inline Base.$(g)(f::typeof($f), x::TrackedArray, y::TrackedReal) = $(g)(SkipOptimize(f), x, y) - @inline Base.$(g)(f::typeof($f), x::TrackedReal, y::TrackedArray) = $(g)(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::TrackedArray, y::TrackedArray) = map(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::TrackedArray, y::TrackedReal) = map(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::TrackedReal, y::TrackedArray) = map(SkipOptimize(f), x, y) end for A in ARRAY_TYPES @eval begin - @inline Base.$(g)(f::typeof($f), x::$A, y::TrackedArray) = $(g)(SkipOptimize(f), x, y) - @inline Base.$(g)(f::typeof($f), x::TrackedArray, y::$A) = $(g)(SkipOptimize(f), x, y) - @inline Base.$(g)(f::typeof($f), x::$A, y::TrackedReal) = $(g)(SkipOptimize(f), x, y) - @inline Base.$(g)(f::typeof($f), x::TrackedReal, y::$A) = $(g)(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::$A, y::TrackedArray) = map(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::TrackedArray, y::$A) = map(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::$A, y::TrackedReal) = map(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::TrackedReal, y::$A) = map(SkipOptimize(f), x, y) end end for R in REAL_TYPES @eval begin - @inline Base.$(g)(f::typeof($f), x::$R, y::TrackedArray) = $(g)(SkipOptimize(f), x, y) - @inline Base.$(g)(f::typeof($f), x::TrackedArray, y::$R) = $(g)(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::$R, y::TrackedArray) = map(SkipOptimize(f), x, y) + @inline Base.map(f::typeof($f), x::TrackedArray, y::$R) = map(SkipOptimize(f), x, y) end end end @@ -34,29 +34,23 @@ end # record # #--------# -for g in (:map,) - @eval @inline Base.$(g)(f::SkipOptimize{F}, t::TrackedArray) where {F} = $(g)(f.f, value(t)) -end +@inline Base.map(f::SkipOptimize{F}, t::TrackedArray) where {F} = map(f.f, value(t)) -for g in (:map,) +@inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) +@inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::TrackedReal) where {F} = map(f.f, value(x), value(y)) +@inline Base.map(f::SkipOptimize{F}, x::TrackedReal, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) +for A in ARRAY_TYPES @eval begin - @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedArray, y::TrackedArray) where {F} = $(g)(f.f, value(x), value(y)) - @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedArray, y::TrackedReal) where {F} = $(g)(f.f, value(x), value(y)) - @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedReal, y::TrackedArray) where {F} = $(g)(f.f, value(x), value(y)) - end - for A in ARRAY_TYPES - @eval begin - @inline Base.$(g)(f::SkipOptimize{F}, x::$A, y::TrackedArray) where {F} = $(g)(f.f, value(x), value(y)) - @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedArray, y::$A) where {F} = $(g)(f.f, value(x), value(y)) - @inline Base.$(g)(f::SkipOptimize{F}, x::$A, y::TrackedReal) where {F} = $(g)(f.f, value(x), value(y)) - @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedReal, y::$A) where {F} = $(g)(f.f, value(x), value(y)) - end + @inline Base.map(f::SkipOptimize{F}, x::$A, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) + @inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::$A) where {F} = map(f.f, value(x), value(y)) + @inline Base.map(f::SkipOptimize{F}, x::$A, y::TrackedReal) where {F} = map(f.f, value(x), value(y)) + @inline Base.map(f::SkipOptimize{F}, x::TrackedReal, y::$A) where {F} = map(f.f, value(x), value(y)) end - for R in REAL_TYPES - @eval begin - @inline Base.$(g)(f::SkipOptimize{F}, x::$R, y::TrackedArray) where {F} = $(g)(f.f, value(x), value(y)) - @inline Base.$(g)(f::SkipOptimize{F}, x::TrackedArray, y::$R) where {F} = $(g)(f.f, value(x), value(y)) - end +end +for R in REAL_TYPES + @eval begin + @inline Base.map(f::SkipOptimize{F}, x::$R, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) + @inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::$R) where {F} = map(f.f, value(x), value(y)) end end @@ -67,26 +61,26 @@ end # dispatch # #----------# -for g in (:map,), (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) +for (M, f, arity) in DiffRules.diffrules(; filter_modules=nothing) if !(isdefined(@__MODULE__, M) && isdefined(getfield(@__MODULE__, M), f)) @warn "$M.$f is not available and hence rule for it can not be defined" continue # Skip rules for methods not defined in the current scope end if arity == 1 - @eval @inline Base.$(g)(f::typeof($M.$f), t::TrackedArray) = $(g)(ForwardOptimize(f), t) + @eval @inline Base.map(f::typeof($M.$f), t::TrackedArray) = map(ForwardOptimize(f), t) elseif arity == 2 (M, f) in SKIPPED_DIFFRULES && continue - @eval @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) + @eval @inline Base.map(f::typeof($M.$f), x::TrackedArray, y::TrackedArray) = map(ForwardOptimize(f), x, y) for A in ARRAY_TYPES @eval begin - @inline Base.$(g)(f::typeof($M.$f), x::$A, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) - @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::$A) = $(g)(ForwardOptimize(f), x, y) + @inline Base.map(f::typeof($M.$f), x::$A, y::TrackedArray) = map(ForwardOptimize(f), x, y) + @inline Base.map(f::typeof($M.$f), x::TrackedArray, y::$A) = map(ForwardOptimize(f), x, y) end end for R in REAL_TYPES @eval begin - @inline Base.$(g)(f::typeof($M.$f), x::$R, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) - @inline Base.$(g)(f::typeof($M.$f), x::TrackedArray, y::$R) = $(g)(ForwardOptimize(f), x, y) + @inline Base.map(f::typeof($M.$f), x::$R, y::TrackedArray) = map(ForwardOptimize(f), x, y) + @inline Base.map(f::typeof($M.$f), x::TrackedArray, y::$R) = map(ForwardOptimize(f), x, y) end end end @@ -95,81 +89,76 @@ end # record # #--------# -for g in (:map,) - @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedArray{X,D}) where {F,X,D} - T = promote_type(X, D) - result = DiffResults.DiffResult(zero(T), zero(T)) - df = v -> ForwardDiff.derivative!(result, f.f, v) - results = $(g)(df, value(x)) +function Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}) where {F,X,D} + T = promote_type(X, D) + result = DiffResults.DiffResult(zero(T), zero(T)) + df = v -> ForwardDiff.derivative!(result, f.f, v) + results = map(df, value(x)) + tp = tape(x) + out = track(map(DiffResults.value, results), D, tp) + cache = (results, df, index_bound(x, out), nothing) + record!(tp, SpecialInstruction, map, x, out, cache) + return out +end + +for A in ARRAY_TYPES + @eval function Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::$A) where {F,X,D} + result = DiffResults.GradientResult(SVector(zero(X))) + df = (vx, vy) -> let vy=vy + ForwardDiff.gradient!(result, s -> f.f(s[1], vy), SVector(vx)) + end + results = map(df, value(x), value(y)) tp = tape(x) out = track(map(DiffResults.value, results), D, tp) - cache = (results, df, index_bound(x, out), nothing) - record!(tp, SpecialInstruction, $(g), x, out, cache) + cache = (results, df, index_bound(x, out), index_bound(y, out)) + record!(tp, SpecialInstruction, map, (x, y), out, cache) return out end - for A in ARRAY_TYPES - @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::$A) where {F,X,D} - result = DiffResults.GradientResult(SVector(zero(X))) - df = (vx, vy) -> let vy=vy - ForwardDiff.gradient!(result, s -> f.f(s[1], vy), SVector(vx)) - end - results = $(g)(df, value(x), value(y)) - tp = tape(x) - out = track(map(DiffResults.value, results), D, tp) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, $(g), (x, y), out, cache) - return out + @eval function Base.map(f::ForwardOptimize{F}, x::$A, y::TrackedArray{Y,D}) where {F,Y,D} + result = DiffResults.GradientResult(SVector(zero(Y))) + df = (vx, vy) -> let vx=vx + ForwardDiff.gradient!(result, s -> f.f(vx, s[1]), SVector(vy)) end - @eval function Base.$(g)(f::ForwardOptimize{F}, x::$A, y::TrackedArray{Y,D}) where {F,Y,D} - result = DiffResults.GradientResult(SVector(zero(Y))) - df = (vx, vy) -> let vx=vx - ForwardDiff.gradient!(result, s -> f.f(vx, s[1]), SVector(vy)) - end - results = $(g)(df, value(x), value(y)) - tp = tape(y) - out = track(map(DiffResults.value, results), D, tp) - cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, $(g), (x, y), out, cache) - return out - end - end - - @eval function Base.$(g)(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {F,X,Y,D} - result = DiffResults.GradientResult(SVector(zero(D), zero(D))) - df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) - results = $(g)(df, value(x), value(y)) - tp = tape(x, y) + results = map(df, value(x), value(y)) + tp = tape(y) out = track(map(DiffResults.value, results), D, tp) cache = (results, df, index_bound(x, out), index_bound(y, out)) - record!(tp, SpecialInstruction, $(g), (x, y), out, cache) + record!(tp, SpecialInstruction, map, (x, y), out, cache) return out end end +function Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {F,X,Y,D} + result = DiffResults.GradientResult(SVector(zero(D), zero(D))) + df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) + results = map(df, value(x), value(y)) + tp = tape(x, y) + out = track(map(DiffResults.value, results), D, tp) + cache = (results, df, index_bound(x, out), index_bound(y, out)) + record!(tp, SpecialInstruction, map, (x, y), out, cache) + return out +end + ################ # forward pass # ################ -for (g!, g) in ((:map!, :map),) - @eval begin - @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof($g)}) - input, output = instruction.input, instruction.output - results, df, _, _ = instruction.cache - if istracked(input) - ($g!)(df, results, value(input)) - else - a, b = input - pull_value!(a) - pull_value!(b) - ($g!)(df, results, value(a), value(b)) - end - output_value = value(output) - for i in eachindex(output_value) - output_value[i] = DiffResults.value(results[i]) - end - return nothing - end +@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(map)}) + input, output = instruction.input, instruction.output + results, df, _, _ = instruction.cache + if istracked(input) + map!(df, results, value(input)) + else + a, b = input + pull_value!(a) + pull_value!(b) + map!(df, results, value(a), value(b)) + end + output_value = value(output) + for i in eachindex(output_value) + output_value[i] = DiffResults.value(results[i]) end + return nothing end ################ diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index c4aa5ef6..90f40815 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -6,7 +6,7 @@ amount scaled by the output's derivative(s). Sometimes, extra partial informatio form of normal scalars/arrays, or `Dual` numbers) needs to accounted for as well. Often, a function's input and output values are not similarly shaped. To account for these -cases, `diffresult` and `reduce` versions of the propagation functions have been implemented. +cases, `diffresult` and `contract` versions of the propagation functions have been implemented. Accumulations into a scalar derivative go through `sum`, whose pairwise reduction is both more accurate and faster than a sequential loop. diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index f14654b2..e0d4a60a 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -365,6 +365,46 @@ end @test isempty(tp) end +@testset "tuple arguments" begin + a = [1.0, 2.0, 3.0] + c = (4.0, 5.0, 6.0) + + tp = InstructionTape() + y = track(copy(a), tp) .* c + @test y isa TrackedArray + @test length(tp) == 1 + @test tp[1].func === ReverseDiff.∇broadcast + @test ReverseDiff.gradient(x -> sum(x .* c), a) == collect(c) + + tape = ReverseDiff.GradientTape(x -> sum(exp.(x .* c)), a) + @test ReverseDiff.gradient!(tape, a) ≈ collect(c) .* exp.(a .* c) + + # a tracked element is not seeded, so it takes the scalar rules + f = x -> sum(x .* (x[1], 5.0, 6.0)) + @test ReverseDiff.gradient(f, a) ≈ ForwardDiff.gradient(f, a) + + # a tuple's elements are checked in turn, including a nested tuple's + h = x -> sum(x .* first.(((x[1], 1.0), (2.0, 3.0), (4.0, 5.0)))) + @test ReverseDiff.gradient(h, a) ≈ ForwardDiff.gradient(h, a) + + tp = InstructionTape() + track(copy(a), tp) .* first.(((1.0, 2.0), (3.0, 4.0), (5.0, 6.0))) + @test length(tp) == 1 +end + +@testset "a `TrackedReal` without a tape is a constant" begin + a = [1.0, 2.0] + c = TrackedReal(2.0, 0.0) + + for (f, expected) in ((x -> x .* c, [2.0, 2.0]), (x -> exp.(x .* c), 2 .* exp.(2 .* a))) + tp = InstructionTape() + x = track(copy(a), tp) + ReverseDiff.seed!(sum(f(x))) + ReverseDiff.reverse_pass!(tp) + @test deriv(x) ≈ expected + end +end + @testset "`NotTracked`" begin f = ReverseDiff.NotTracked(t -> 2t) @@ -382,8 +422,8 @@ end return @allocated f(targs...) end - # the output's value and derivative take 8 B per element each, and one `Dual` per - # element would add at least 16 B more + # bytes allocated by the whole recording: the output's value and derivative take 8 B per + # element each, and one `Dual` per element would add at least 16 B more for f in (t -> t .* 2.0, t -> 2.0 .* t, t -> t .+ 1.0, t -> 1.0 .- t, t -> t ./ 4.0, t -> 4.0 .\ t, t -> identity.(t), t -> .-t) @test recordbytes(f, v) < 24n diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index 53add506..b981baf4 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -52,8 +52,10 @@ function test_elementwise(f, fopt, x, tp) # record yt = broadcast(fopt, xt) @test yt == y - # a function constant in its argument has no derivative and is left off the tape - @test length(tp) == (yt isa ReverseDiff.TrackedArray ? 1 : 0) + # a function that returns no `Dual` has no derivative and is left off the tape + tracked = f(ForwardDiff.Dual(first(x), 1.0)) isa ForwardDiff.Dual + @test (yt isa ReverseDiff.TrackedArray) == tracked + @test length(tp) == tracked # reverse out = similar(y, (length(x), length(x))) From d20947f1b4b2c4ce017ed332e5f6efe564442677 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 28 Sep 2026 15:12:06 +0200 Subject: [PATCH 19/36] Unwrap structured matrices by dispatch instead of stripping `Dual`s in `value` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `value(::Dual)` for `∇broadcast`'s tag was global, so a broadcast nested in another stripped the outer perturbation and returned a zero gradient. `value` now unwraps `LinearAlgebra`'s wrappers of a `TrackedArray` by dispatch, since `map` over a `StructuredMatrix` broadcasts, and a nested broadcast raises the descriptive error again. This also replaces `mulargvalue`. The broadcast tag is keyed on the seeded arguments' value type, as ForwardDiff keys it, rather than on the derivative type, and `D` is computed where the output is tracked. Partials are numbered left to right, and known partials unwrap only the arguments they name. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 54 ++++++++++++---------------- src/derivatives/linalg/arithmetic.jl | 12 +++---- src/derivatives/propagation.jl | 4 +-- src/tracked.jl | 12 +++++++ test/derivatives/BroadcastTests.jl | 35 ++++++++++-------- 5 files changed, 62 insertions(+), 55 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 2ac5eac3..00063560 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -143,17 +143,15 @@ end @inline incr(::Val{k}) where {k} = Val(k + 1) # argument `i` has partial `slots[i]` (`Val(0)` if untracked), partial `k` belongs to argument -# `positions[k]`, numbered from the last tracked argument. Built on the way out of the recursion -# to stay constant-folded. +# `positions[k]`. Built on the way out of the recursion to stay constant-folded. @inline trackedslots(::Tuple{}) = (), (), Val(0) @inline function trackedslots(args::Tuple) - slots, positions, n = trackedslots(Base.tail(args)) - positions = map(incr, positions) - if istracked(first(args)) + slots, positions, n = trackedslots(Base.front(args)) + if istracked(last(args)) k = incr(n) - return (k, slots...), (positions..., Val(1)), k + return (slots..., k), (positions..., Val(length(args))), k else - return (Val(0), slots...), positions, n + return (slots..., Val(0)), positions, n end end @@ -161,6 +159,9 @@ end # `DiffRules` leaves partials such as the order of `besselj` as `NaN`, which `NaN * 0` # would spread to every slot, so untracked arguments are not dualized +@inline dualvaltype(::Val{0}, v) = Union{} +@inline dualvaltype(::Val, v) = eltype(v) + @inline dualize(::Type, ::Val{0}, ::Val, x) = x @inline function dualize(::Type{T}, ::Val{k}, valP::Val, x) where {T, k} return Dual{T}(x, ntuple(j -> j == k, valP)) @@ -174,7 +175,7 @@ end const RealOrArray = Union{Real, AbstractArray{<:Real}} # an argument read at the output's own index has to span the whole output -ifsameshape(c::Contract, x, y) = x isa Real || y isa Real || axes(x) == axes(y) ? c : nothing +ifsameshape(c::Contract, x, y) = (x isa Real || y isa Real || axes(x) == axes(y)) ? c : nothing # partial with respect to argument `i`, or `nothing`. Requiring every argument to be # `RealOrArray` makes positions in `args` positions in the arguments `splitargs` keeps. @@ -211,21 +212,14 @@ function knownpartials(f, args, positions::Tuple) end end -# marks the `Dual`s seeded by `∇broadcast`, so `value` can strip them as it strips tracking; -# `F` gives each function its own tag, which `ForwardDiff` orders by first use, as `Tag(f, V)` would -struct BroadcastTag{F} end - -value(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = ForwardDiff.value(T, x) - # at least one argument has to be a non-0-dimensional array: `copy` sends the scalar and # 0-dimensional cases onto the scalar rules instead @inline function ∇broadcast(f::F, args::Tuple, argvals::Tuple) where {F} inds, targs, untracked = splitargs(args) _, vals, _ = splitargs(argvals) - D = mapreduce(getouttype, promote_type, targs) slots, positions, valP = trackedslots(targs) # one tag per broadcast, and a new one under an enclosing differentiation - T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), D)) + T = typeof(ForwardDiff.Tag(f, reduce(promote_type, map(dualvaltype, slots, vals)))) # `broadcast` calls `df` elementwise, so it receives one scalar per argument function df(x::Vararg{Any,N}) where {N} dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) @@ -235,11 +229,10 @@ value(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = ForwardDiff.value vf(x::Vararg{Any,N}) where {N} = splatcall(f, x, untracked, inds) entries = knownpartials(f, args, positions) if entries === nothing - results = broadcast(df, vals...) + return trackresults(T, broadcast(df, vals...), df, vf, targs, vals) else - results = KnownPartials(entries) + return recordresults(T, KnownPartials(entries), df, vf, targs, vals) end - return trackresults(T, results, df, vf, targs, vals, D) end # the cache carries what the replay needs: `df` to recompute the stored partials, or, where @@ -262,8 +255,8 @@ replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) return nothing end -@inline function recordresults(::Type{T}, results, df, vf, targs, vals, - ::Type{D}) where {T, D} +@inline function recordresults(::Type{T}, results, df, vf, targs, vals) where {T} + D = mapreduce(getouttype, promote_type, targs) foreach(v -> checkargtags(D, eltype(v)), vals) g, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) @@ -275,9 +268,6 @@ end return out end -@inline trackresults(::Type{T}, results::KnownPartials, df, vf, targs, vals, - ::Type{D}) where {T, D} = recordresults(T, results, df, vf, targs, vals, D) - # an enclosing differentiation's tag is constant in our arguments, while one nested inside # `f` buries our partial; a widened element type is a `Union`, so every member is checked @inline function checktags(::Type{T}, ::Type{E}) where {T, E} @@ -295,13 +285,12 @@ end # a type-unstable `f`, such as one with an integer literal branch, leaves an abstract # element type that can still hide a `Dual{T}` -@inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, vals, - ::Type{D}) where {T, D} +@inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, vals) where {T} if typeintersect(eltype(results), Dual{T}) === Union{} checktags(T, eltype(results)) return results else - return recordresults(T, results, df, vf, targs, vals, D) + return recordresults(T, results, df, vf, targs, vals) end end @@ -325,10 +314,13 @@ struct PartialsWithArgs{E<:Tuple,A<:Tuple} args::A end -# the partials of one reverse pass. Only known partials read the argument values, so the -# per-element cache skips computing them, which copies an array of `TrackedReal`s. +# the partials of one reverse pass. Only known partials read argument values, and only of the +# arguments they name, since unwrapping an array of `TrackedReal`s copies it. reversepartials(results::AbstractArray, _) = results -reversepartials(p::KnownPartials, input) = PartialsWithArgs(p.entries, map(value, input)) +function reversepartials(p::KnownPartials, input) + args = map(e -> map(j -> value(getat(input, j)), e.args), p.entries) + return PartialsWithArgs(p.entries, args) +end # an argument broadcast to the full output shape needs no index clamping function _br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, @@ -347,7 +339,7 @@ _br_add_to_deriv!(::Type{T}, x, slot::Val, out_deriv, results, ::Nothing) where _increment_deriv!(::Type{T}, x, out_deriv, results, ::Val{k}, bound...) where {T, k} = diffresult_increment_deriv!(T, x, out_deriv, results, k, bound...) _increment_deriv!(::Type, x, out_deriv, p::PartialsWithArgs, ::Val{k}, bound...) where {k} = - contract_increment_deriv!(x, out_deriv, p.entries[k], p.args, bound...) + contract_increment_deriv!(x, out_deriv, p.entries[k], p.args[k], bound...) @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) input, output = instruction.input, instruction.output diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index b3cb4b73..d0c15118 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -157,10 +157,6 @@ end # multiplication (*) # ###################### -mulargvalue(x) = value(x) -mulargvalue(x::Adjoint) = adjoint(value(adjoint(x))) -mulargvalue(x::Transpose) = transpose(value(transpose(x))) - mulargpullvalue!(x) = pull_value!(x) mulargpullvalue!(x::Adjoint) = pull_value!(adjoint(x)) mulargpullvalue!(x::Transpose) = pull_value!(transpose(x)) @@ -191,14 +187,14 @@ end @inline function record_mul(x, y, ::Type{D}) where D tp = tape(x, y) - out = track(*(mulargvalue(x), mulargvalue(y)), D, tp) + out = track(*(value(x), value(y)), D, tp) cache = (similar(x, D), similar(y, D)) record!(tp, SpecialInstruction, *, (x, y), out, cache) return out end @inline function record_mul!(out::TrackedArray{V,D}, x, y) where {V,D} - copyto!(mulargvalue(out), *(mulargvalue(x), mulargvalue(y))) + copyto!(value(out), *(value(x), value(y))) cache = (similar(x, D), similar(y, D)) record!(tape(x, y), SpecialInstruction, *, (x, y), out, cache) return out @@ -281,9 +277,9 @@ end mulargpullvalue!(a) mulargpullvalue!(b) if instruction.output isa Number - value!(instruction.output, mulargvalue(a) * mulargvalue(b)) + value!(instruction.output, value(a) * value(b)) else - mul!(mulargvalue(instruction.output), mulargvalue(a), mulargvalue(b)) + mul!(value(instruction.output), value(a), value(b)) end return nothing end diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index 90f40815..7da73b8e 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -108,11 +108,11 @@ struct Contract{Op,A<:Tuple} args::A end -_at(::Val{j}, i, args) where {j} = _elem(args[j], i) _elem(v::Real, i) = v _elem(v::AbstractArray, i) = v[i] -_contract(e::Contract, seed, i, args) = e.op(seed, map(a -> _at(a, i, args), e.args)...) +# `vals` holds the values of the arguments `e` names +_contract(e::Contract, seed, i, vals) = e.op(seed, map(v -> _elem(v, i), vals)...) function diffresult_increment_deriv!(::Type{T}, input::AbstractArray, x::AbstractArray, results::AbstractArray, p::Int) where {T} diff --git a/src/tracked.jl b/src/tracked.jl index 31c9f345..83775760 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -106,6 +106,18 @@ end @inline value(x::AbstractArray) = istracked(x) ? map(value, x) : x @inline value(t::TrackedReal) = t.value @inline value(t::TrackedArray) = t.value +value(A::Adjoint) = adjoint(value(parent(A))) +value(A::Transpose) = transpose(value(parent(A))) +value(D::Diagonal) = Diagonal(value(D.diag)) +value(B::Bidiagonal) = Bidiagonal(value(B.dv), value(B.ev), Symbol(B.uplo)) +value(T::Tridiagonal) = Tridiagonal(value(T.dl), value(T.d), value(T.du)) +value(S::SymTridiagonal) = SymTridiagonal(value(S.dv), value(S.ev)) +value(S::Symmetric) = Symmetric(value(parent(S)), Symbol(S.uplo)) +value(H::Hermitian) = Hermitian(value(parent(H)), Symbol(H.uplo)) +value(H::UpperHessenberg) = UpperHessenberg(value(parent(H))) +for W in (:UpperTriangular, :LowerTriangular, :UnitUpperTriangular, :UnitLowerTriangular) + @eval value(A::$W) = $W(value(parent(A))) +end @inline deriv(t::TrackedArray) = t.deriv @inline deriv(t::TrackedReal) = t.deriv diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index e0d4a60a..eac4e2c6 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -261,10 +261,10 @@ end 1.0, ForwardDiff.Dual{tagA}(2.0, 3.0)] nested = Union{Float64,ForwardDiff.Dual{tagB,Float64,1}}[1.0] - @test ReverseDiff.trackresults(tagB, concrete, identity, identity, (), (), Float64) === concrete - @test ReverseDiff.trackresults(tagB, widened, identity, identity, (), (), Float64) === widened + @test ReverseDiff.trackresults(tagB, concrete, identity, identity, (), ()) === concrete + @test ReverseDiff.trackresults(tagB, widened, identity, identity, (), ()) === widened @test_throws ForwardDiff.DualMismatchError ReverseDiff.trackresults( - tagA, nested, identity, identity, (), (), Float64) + tagA, nested, identity, identity, (), ()) @test ReverseDiff.getpartial(tagA, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 2.0 @test ReverseDiff.getpartial(tagA, 1.0, 1) == 0.0 @@ -290,6 +290,10 @@ end # a tape whose derivatives are themselves `Dual`s can hold it, so it is left alone h(a) = ReverseDiff.gradient(x -> sum(x .* a), [ForwardDiff.Dual(1.0, 0.0)]) @test h(3.0) == [ForwardDiff.Dual(3.0, 0.0)] + + # also when the perturbation comes from an enclosing broadcast + f(x) = sum(broadcast(a -> ReverseDiff.gradient(y -> sum(y .* a), [1.0])[1], x)) + @test_throws ArgumentError ReverseDiff.gradient(f, [2.0, 3.0]) end @testset "zero-dimensional arrays (#265)" begin @@ -348,21 +352,24 @@ end @test gm == Diagonal(a) end -@testset "`value` in a broadcast is not recorded" begin - a = rand(3) +@testset "`value` of a wrapped `TrackedArray` is not recorded" begin + wrap(d, e, A) = (Diagonal(d), Bidiagonal(d, e, :L), Tridiagonal(e, d, e), SymTridiagonal(d, e), + UpperTriangular(A), LowerTriangular(A), UnitUpperTriangular(A), + UnitLowerTriangular(A), UpperHessenberg(A), Symmetric(A, :L), + Hermitian(A, :L), adjoint(A), transpose(A)) + d, e, A = rand(3), rand(2), rand(3, 3) tp = InstructionTape() - x = track(copy(a), tp) + expected = wrap(d, e, A) + actual = map(value, wrap(track(d, tp), track(e, tp), track(A, tp))) - d = value(Diagonal(x)) - @test d isa Diagonal{Float64,Vector{Float64}} - @test d == Diagonal(a) + @test all(map((a, b) -> a isa typeof(b), actual, expected)) + @test actual == expected @test isempty(tp) - # also when fused with a differentiable function - y = value.(x .+ 1) - @test y isa Vector{Float64} - @test y ≈ a .+ 1 - @test isempty(tp) + # the test covers every `StructuredMatrix` type + covered = map(nameof ∘ typeof, expected) + structured = Base.uniontypes(Base.unwrap_unionall(LinearAlgebra.StructuredMatrix)) + @test issubset(map(nameof, structured), covered) end @testset "tuple arguments" begin From 2146fbbba4bf09c77f9db94fe409c1c1b217817a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 28 Sep 2026 15:38:34 +0200 Subject: [PATCH 20/36] Differentiate `map` w.r.t. an array of `TrackedReal`s next to a `TrackedArray` A `ForwardOptimize` `map` over a `TrackedArray` and a plain array treated the latter as constant, yet the reverse pass read its partial whenever its element type admits `TrackedReal`, so `map(hypot, x, Real[...])` threw a `BoundsError`, on master too. Such an array is now differentiated like a second `TrackedArray`. The `SkipOptimize` `map` methods taking a `TrackedReal` go: `Base.map` and the scalar `@skip` rules give the same results and leave the tape empty. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/elementwise.jl | 20 +++++++++++--------- test/derivatives/ElementwiseTests.jl | 17 +++++++++++++++++ 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/src/derivatives/elementwise.jl b/src/derivatives/elementwise.jl index 35c58f0a..76fd3efa 100644 --- a/src/derivatives/elementwise.jl +++ b/src/derivatives/elementwise.jl @@ -12,15 +12,11 @@ end for f in SKIPPED_BINARY_SCALAR_FUNCS @eval begin @inline Base.map(f::typeof($f), x::TrackedArray, y::TrackedArray) = map(SkipOptimize(f), x, y) - @inline Base.map(f::typeof($f), x::TrackedArray, y::TrackedReal) = map(SkipOptimize(f), x, y) - @inline Base.map(f::typeof($f), x::TrackedReal, y::TrackedArray) = map(SkipOptimize(f), x, y) end for A in ARRAY_TYPES @eval begin @inline Base.map(f::typeof($f), x::$A, y::TrackedArray) = map(SkipOptimize(f), x, y) @inline Base.map(f::typeof($f), x::TrackedArray, y::$A) = map(SkipOptimize(f), x, y) - @inline Base.map(f::typeof($f), x::$A, y::TrackedReal) = map(SkipOptimize(f), x, y) - @inline Base.map(f::typeof($f), x::TrackedReal, y::$A) = map(SkipOptimize(f), x, y) end end for R in REAL_TYPES @@ -37,14 +33,10 @@ end @inline Base.map(f::SkipOptimize{F}, t::TrackedArray) where {F} = map(f.f, value(t)) @inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) -@inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::TrackedReal) where {F} = map(f.f, value(x), value(y)) -@inline Base.map(f::SkipOptimize{F}, x::TrackedReal, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) for A in ARRAY_TYPES @eval begin @inline Base.map(f::SkipOptimize{F}, x::$A, y::TrackedArray) where {F} = map(f.f, value(x), value(y)) @inline Base.map(f::SkipOptimize{F}, x::TrackedArray, y::$A) where {F} = map(f.f, value(x), value(y)) - @inline Base.map(f::SkipOptimize{F}, x::$A, y::TrackedReal) where {F} = map(f.f, value(x), value(y)) - @inline Base.map(f::SkipOptimize{F}, x::TrackedReal, y::$A) where {F} = map(f.f, value(x), value(y)) end end for R in REAL_TYPES @@ -101,8 +93,12 @@ function Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}) where {F,X,D} return out end +# an array of `TrackedReal`s is differentiated like a `TrackedArray` for A in ARRAY_TYPES @eval function Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::$A) where {F,X,D} + if istracked(y) + return record_map(f, x, y, D) + end result = DiffResults.GradientResult(SVector(zero(X))) df = (vx, vy) -> let vy=vy ForwardDiff.gradient!(result, s -> f.f(s[1], vy), SVector(vx)) @@ -115,6 +111,9 @@ for A in ARRAY_TYPES return out end @eval function Base.map(f::ForwardOptimize{F}, x::$A, y::TrackedArray{Y,D}) where {F,Y,D} + if istracked(x) + return record_map(f, x, y, D) + end result = DiffResults.GradientResult(SVector(zero(Y))) df = (vx, vy) -> let vx=vx ForwardDiff.gradient!(result, s -> f.f(vx, s[1]), SVector(vy)) @@ -128,7 +127,10 @@ for A in ARRAY_TYPES end end -function Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {F,X,Y,D} +Base.map(f::ForwardOptimize{F}, x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {F,X,Y,D} = + record_map(f, x, y, D) + +function record_map(f::ForwardOptimize, x, y, ::Type{D}) where {D} result = DiffResults.GradientResult(SVector(zero(D), zero(D))) df = (vx, vy) -> ForwardDiff.gradient!(result, s -> f.f(s[1], s[2]), SVector(vx, vy)) results = map(df, value(x), value(y)) diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index b981baf4..62e70640 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -486,4 +486,21 @@ end end end +@testset "`map` over a `TrackedArray` and an array that may hold `TrackedReal`s" begin + c = [1.0, 2.0, 3.0] + ∇hypot(a, b) = [a ./ hypot.(a, b); b ./ hypot.(a, b)] + + cases = ((t -> sum(map(hypot, t[1:3], [t[4], t[5], t[6]])), t -> ∇hypot(t[1:3], t[4:6])), + (t -> sum(map(hypot, [t[1], t[2], t[3]], t[4:6])), t -> ∇hypot(t[1:3], t[4:6])), + (t -> sum(map(hypot, t[1:3], Real[c...])), t -> [∇hypot(t[1:3], c)[1:3]; zeros(3)])) + for (f, ∇f) in cases + v = rand(6) + @test ReverseDiff.gradient(f, v) ≈ ∇f(v) + + tape = ReverseDiff.GradientTape(f, v) + w = rand(6) + @test ReverseDiff.gradient!(tape, w) ≈ ∇f(w) + end +end + end # module From 6581b2a0e9e6244043783c1a907dbb3dd7b73d4a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 28 Sep 2026 20:02:05 +0200 Subject: [PATCH 21/36] Drop broadcast perturbations in `@skip` only, and unwrap only `Diagonal` by hand MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `value` stopped stripping `∇broadcast`'s `Dual`s, which `@skip` relied on: a fused `x .* @skip(exp).(x)` differentiated `exp`. The tag wraps `BroadcastTag{F}` again and `skipvalue` drops it, and nested tracking, at every order. The rewrapping `value` methods mapped over whole parents, which LinearAlgebra leaves `#undef` outside a triangle. The generic fallback handles every wrapper but `Diagonal`, whose `map` broadcasts and would be recorded. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 7 ++++++- src/macros.jl | 8 ++++---- src/tracked.jl | 9 --------- test/derivatives/BroadcastTests.jl | 7 ++++++- 4 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 00063560..0cb55d74 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -212,6 +212,11 @@ function knownpartials(f, args, positions::Tuple) end end +# marks the `Dual`s `∇broadcast` seeds, so `@skip` can drop them as it drops tracking +struct BroadcastTag{F} end + +skipvalue(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = skipvalue(ForwardDiff.value(T, x)) + # at least one argument has to be a non-0-dimensional array: `copy` sends the scalar and # 0-dimensional cases onto the scalar rules instead @inline function ∇broadcast(f::F, args::Tuple, argvals::Tuple) where {F} @@ -219,7 +224,7 @@ end _, vals, _ = splitargs(argvals) slots, positions, valP = trackedslots(targs) # one tag per broadcast, and a new one under an enclosing differentiation - T = typeof(ForwardDiff.Tag(f, reduce(promote_type, map(dualvaltype, slots, vals)))) + T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), reduce(promote_type, map(dualvaltype, slots, vals)))) # `broadcast` calls `df` elementwise, so it receives one scalar per argument function df(x::Vararg{Any,N}) where {N} dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) diff --git a/src/macros.jl b/src/macros.jl index 5e15acd1..1be3df46 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -162,10 +162,10 @@ macro skip(ex) return esc(annotate_func_expr(:SkipOptimize, __module__, ex)) end -@inline (self::SkipOptimize{F})(args...) where {F} = self.f(map(value, args)...) -@inline (self::SkipOptimize{F})(a) where {F} = self.f(value(a)) -@inline (self::SkipOptimize{F})(a, b) where {F} = self.f(value(a), value(b)) -@inline (self::SkipOptimize{F})(a, b, c) where {F} = self.f(value(a), value(b), value(c)) +skipvalue(x) = value(x) +skipvalue(x::TrackedReal) = skipvalue(value(x)) + +@inline (self::SkipOptimize{F})(args...) where {F} = self.f(map(skipvalue, args)...) """ f(x) = dot(x, x) diff --git a/src/tracked.jl b/src/tracked.jl index 83775760..5b344127 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -109,15 +109,6 @@ end value(A::Adjoint) = adjoint(value(parent(A))) value(A::Transpose) = transpose(value(parent(A))) value(D::Diagonal) = Diagonal(value(D.diag)) -value(B::Bidiagonal) = Bidiagonal(value(B.dv), value(B.ev), Symbol(B.uplo)) -value(T::Tridiagonal) = Tridiagonal(value(T.dl), value(T.d), value(T.du)) -value(S::SymTridiagonal) = SymTridiagonal(value(S.dv), value(S.ev)) -value(S::Symmetric) = Symmetric(value(parent(S)), Symbol(S.uplo)) -value(H::Hermitian) = Hermitian(value(parent(H)), Symbol(H.uplo)) -value(H::UpperHessenberg) = UpperHessenberg(value(parent(H))) -for W in (:UpperTriangular, :LowerTriangular, :UnitUpperTriangular, :UnitLowerTriangular) - @eval value(A::$W) = $W(value(parent(A))) -end @inline deriv(t::TrackedArray) = t.deriv @inline deriv(t::TrackedReal) = t.deriv diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index eac4e2c6..0b54b682 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -124,6 +124,12 @@ end @test y isa Matrix{Float64} @test y == exp.(value(x)) @test isempty(tp) + + # fused, and at every order + f(v) = sum(v .* ReverseDiff.@skip(exp).(v)) + a = rand(3) + @test ReverseDiff.gradient(f, a) ≈ exp.(a) + @test ReverseDiff.hessian(f, a) ≈ zeros(3, 3) end @testset "a type broadcast as a function does not force the fallback path" begin @@ -362,7 +368,6 @@ end expected = wrap(d, e, A) actual = map(value, wrap(track(d, tp), track(e, tp), track(A, tp))) - @test all(map((a, b) -> a isa typeof(b), actual, expected)) @test actual == expected @test isempty(tp) From 13e2096c0526b83160697e27af6cc78dc16b7f1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 29 Sep 2026 01:13:56 +0200 Subject: [PATCH 22/36] Record broadcasts by inferred type, keep static fallbacks static, and reject buried tracked values Whether a broadcast was recorded depended on its recorded values: a branchy `f` whose every element took a constant branch, such as `ifelse.(x .> 0, x, 0.0)` at negative `x`, was left off the tape, and replays gave wrong gradients. The decision now uses the element type `combine_eltypes` infers, and the results are converted to it so that a replay can store the other branch. The broadcast tag is keyed on every argument's type, so a `Dual` of an enclosing differentiation on an untracked argument always makes a newer tag; with an older one, forward-over-reverse threw "Cannot determine ordering". Arguments with an abstract element type are checked element by element. The fallback no longer copies values to find the untracked style, and a `TrackedArray` of a static array is broadcast as a static array of its elements, which crashed before. `@skip` drops every level of tracking from arrays too. A `Dual` wrapping a `TrackedReal` is a perturbation of an enclosing ForwardDiff call that reached the tape, and was differentiated as a constant. Such outputs now raise an error, promotion keeps them apart from `TrackedReal`s rather than cutting them off the tape, and converting one to a `TrackedReal` raises an error. Reverse-over-forward with such mixed arrays now works. Co-Authored-By: Claude Opus 5.5 --- src/api/utils.jl | 1 + src/derivatives/broadcast.jl | 75 ++++++++++++++++------------- src/macros.jl | 1 + src/tracked.jl | 17 ++++++- test/derivatives/BroadcastTests.jl | 76 ++++++++++++++++++++++++++---- 5 files changed, 129 insertions(+), 41 deletions(-) diff --git a/src/api/utils.jl b/src/api/utils.jl index 147cbd80..38028be3 100644 --- a/src/api/utils.jl +++ b/src/api/utils.jl @@ -34,6 +34,7 @@ function seeded_reverse_pass!(result, output::TrackedReal, input, tape) end function seeded_reverse_pass!(result, output::Number, input, tape) + checkoutput(output) result = extract_result!(result, output, input) return result end diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 0cb55d74..ea6e885f 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -64,10 +64,17 @@ Broadcast.BroadcastStyle(::TrackedStyle{M}, ::DefaultArrayStyle{N}) where {M,N} Broadcast.BroadcastStyle(::TrackedStyle{N}, ::DefaultArrayStyle{N}) where {N} = TrackedStyle{N}() Broadcast.BroadcastStyle(::TrackedStyle{Any}, ::DefaultArrayStyle) = TrackedStyle{Any}() -# the untracked style decides where the fallback re-dispatches, and a `CuArray` backing a -# `TrackedArray` has to keep its own -recur_value(xs) = xs -recur_value(xs::Union{TrackedReal, AbstractArray{<:TrackedReal}}) = recur_value(value(xs)) +# the style the arguments have untracked, which picks the fallback's output container +untrackedstyle(x) = Broadcast.BroadcastStyle(typeof(x)) +untrackedstyle(x::Union{TrackedReal, AbstractArray{<:TrackedReal}}) = untrackedstyle(value(x)) +# unwrapping would copy +untrackedstyle(::Array{<:TrackedReal,N}) where {N} = DefaultArrayStyle{N}() + +# a static value's elements, as a static array that `StaticArrays` broadcasts over itself +staticelements(x) = x +function staticelements(t::TrackedArray{<:Any,<:Any,<:Any,VA}) where {VA<:StaticArray} + return similar_type(VA, eltype(t))(ntuple(i -> t[i], Val(length(VA)))) +end remove_not_tracked(f) = f remove_not_tracked(f::NotTracked) = f.f @@ -83,19 +90,19 @@ function Base.copy(_bc::Broadcasted{<:TrackedStyle}) if axes(_bc) isa Tuple{} return _bc[CartesianIndex()] end - bc = remove_not_tracked(_bc) - flattened_bc = Base.Broadcast.flatten(bc) + flattened_bc = Base.Broadcast.flatten(remove_not_tracked(_bc)) f, args = flattened_bc.f, flattened_bc.args - vals = map(value, args) # only the arguments are seeded, not e.g. a closure's captures, and a `TrackedArray` # holds only `Real`s - if mayhidetracked(_bc) || !(Broadcast.combine_eltypes(f, vals) <: Real) - axs = flattened_bc.axes - style = typeof(Broadcast.combine_styles(map(recur_value, vals)...)) - return copy(Broadcasted{style, typeof(axs), typeof(f), typeof(args)}(f, args, axs)) - else - return ∇broadcast(f, args, vals) + if !mayhidetracked(_bc) + vals = map(value, args) + if Broadcast.combine_eltypes(f, vals) <: Real + return ∇broadcast(f, args, vals) + end end + elargs = map(staticelements, args) + style = typeof(reduce(Broadcast.result_style, map(untrackedstyle, elargs))) + return copy(Broadcast.instantiate(Broadcasted{style}(f, elargs))) end _no_tracked_dest() = throw(ArgumentError("`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead.")) @@ -159,9 +166,6 @@ end # `DiffRules` leaves partials such as the order of `besselj` as `NaN`, which `NaN * 0` # would spread to every slot, so untracked arguments are not dualized -@inline dualvaltype(::Val{0}, v) = Union{} -@inline dualvaltype(::Val, v) = eltype(v) - @inline dualize(::Type, ::Val{0}, ::Val, x) = x @inline function dualize(::Type{T}, ::Val{k}, valP::Val, x) where {T, k} return Dual{T}(x, ntuple(j -> j == k, valP)) @@ -187,11 +191,11 @@ knownpartial(::typeof(-), ::Val{1}, ::Tuple{RealOrArray}) = Contract(-, ()) knownpartial(::typeof(-), ::Val{1}, ::Tuple{RealOrArray,RealOrArray}) = Contract(identity, ()) knownpartial(::typeof(-), ::Val{2}, ::Tuple{RealOrArray,RealOrArray}) = Contract(-, ()) -# a denominator's partial `-x/y^2` is no argument of the broadcast knownpartial(::typeof(*), ::Val{1}, (x, y)::Tuple{RealOrArray,RealOrArray}) = ifsameshape(Contract(*, (Val(2),)), x, y) knownpartial(::typeof(*), ::Val{2}, (x, y)::Tuple{RealOrArray,RealOrArray}) = ifsameshape(Contract(*, (Val(1),)), x, y) +# a denominator's partial `-x/y^2` is no argument of the broadcast knownpartial(::typeof(/), ::Val{1}, (x, y)::Tuple{RealOrArray,RealOrArray}) = ifsameshape(Contract(/, (Val(2),)), x, y) knownpartial(::typeof(\), ::Val{2}, (x, y)::Tuple{RealOrArray,RealOrArray}) = @@ -223,8 +227,8 @@ skipvalue(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = skipvalue(For inds, targs, untracked = splitargs(args) _, vals, _ = splitargs(argvals) slots, positions, valP = trackedslots(targs) - # one tag per broadcast, and a new one under an enclosing differentiation - T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), reduce(promote_type, map(dualvaltype, slots, vals)))) + # keyed on every argument's type, so an enclosing differentiation's `Dual` makes a newer tag + T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), typeof(argvals))) # `broadcast` calls `df` elementwise, so it receives one scalar per argument function df(x::Vararg{Any,N}) where {N} dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) @@ -246,13 +250,9 @@ replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} = (df, map(y -> ForwardDiff.value(T, y), results)) replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) -# a perturbation riding on an argument ends up in the derivative, so `D` has to be able to hold -# it; a widened element type is a `Union` +# a perturbation riding on an argument ends up in the derivative, so `D` has to be able to hold it @inline function checkargtags(::Type{D}, ::Type{E}) where {D, E} - if E isa Union - checkargtags(D, E.a) - checkargtags(D, E.b) - elseif ForwardDiff.tagtype(E) !== Nothing && !(promote_type(D, E) <: D) + if ForwardDiff.tagtype(E) !== Nothing && !(promote_type(D, E) <: D) throw(ArgumentError(LazyString("a broadcast argument with element type ", E, " carries a perturbation that a derivative of type ", D, " cannot hold"))) @@ -260,9 +260,19 @@ replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) return nothing end +# an abstract element type such as `Real` can hide a perturbation +function checkargvalues(::Type{D}, v) where {D} + if isconcretetype(eltype(v)) + checkargtags(D, eltype(v)) + else + foreach(x -> checkargvalues(D, x), v) + end + return nothing +end + @inline function recordresults(::Type{T}, results, df, vf, targs, vals) where {T} D = mapreduce(getouttype, promote_type, targs) - foreach(v -> checkargtags(D, eltype(v)), vals) + foreach(v -> checkargvalues(D, v), vals) g, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) out = track(outvalue, D, tp) @@ -288,14 +298,15 @@ end return nothing end -# a type-unstable `f`, such as one with an integer literal branch, leaves an abstract -# element type that can still hide a `Dual{T}` +# inferred, not read off the results, since a replay can take other branches and writes into +# the results @inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, vals) where {T} - if typeintersect(eltype(results), Dual{T}) === Union{} - checktags(T, eltype(results)) + E = Broadcast.combine_eltypes(df, vals) + if typeintersect(E, Dual{T}) === Union{} + checktags(T, E) return results else - return recordresults(T, results, df, vf, targs, vals) + return recordresults(T, convert(AbstractArray{E}, results), df, vf, targs, vals) end end @@ -306,7 +317,7 @@ end results, _, tag, targets = instruction.cache T = typeof(tag) partials = reversepartials(results, input) - map(targets) do (p, k, bound) + foreach(targets) do (p, k, bound) x = getat(input, p) istracked(x) && _br_add_to_deriv!(T, x, k, output_deriv, partials, bound) end diff --git a/src/macros.jl b/src/macros.jl index 1be3df46..27548b13 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -164,6 +164,7 @@ end skipvalue(x) = value(x) skipvalue(x::TrackedReal) = skipvalue(value(x)) +skipvalue(x::AbstractArray{<:TrackedReal}) = skipvalue(value(x)) @inline (self::SkipOptimize{F})(args...) where {F} = self.f(map(skipvalue, args)...) diff --git a/src/tracked.jl b/src/tracked.jl index 5b344127..c292feb7 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -205,7 +205,14 @@ push_deriv!(x::AbstractArray) = (istracked(x) && foreach(push_deriv!, x); nothin # seed/unseed # #-------------# -seed!(::Real) = nothing +# a `Dual` of an enclosing differentiation buries a tracked value, whose derivative would +# silently be zero +checkoutput(x) = nothing +function checkoutput(::Dual{<:Any,<:TrackedReal}) + throw(ArgumentError("ForwardDiff cannot differentiate through ReverseDiff (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)")) +end + +seed!(x::Real) = checkoutput(x) seed!(t::TrackedReal) = (t.deriv = one(derivtype(t)); push_deriv!(t); nothing) seed!(t::TrackedArray, i) = (t.deriv[i] = one(derivtype(t)); nothing) seed!(x::AbstractArray, i) = seed!(x[i]) @@ -272,6 +279,11 @@ function Base.convert(::Type{T}, x::R) where {T<:TrackedReal, R<:Real} ) end +# wrapping cuts the tracked value off its tape +function Base.convert(::Type{<:TrackedReal}, ::Dual{<:Any,<:TrackedReal}) + throw(ArgumentError("this nesting of ForwardDiff and ReverseDiff is not supported: a `Dual` of tracked numbers cannot be converted to a tracked number (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)")) +end + Base.convert(::Type{T}, t::T) where {T<:TrackedReal} = t Base.convert(::Type{T}, t::T) where {T<:TrackedArray} = t @@ -291,6 +303,9 @@ end # Avoid method ambiguities for ForwardDiff.Dual Base.promote_rule(::Type{TrackedReal{V1,D,O}}, ::Type{Dual{T,V2,N}}) where {V1,D,O,T,V2,N} = TrackedReal{promote_type(V1,Dual{T,V2,N}),D,O} Base.promote_rule(::Type{Dual{T,V1,N}}, ::Type{TrackedReal{V2,D,O}}) where {T,V1,N,V2,D,O} = TrackedReal{promote_type(Dual{T,V1,N},V2),D,O} +# a `Dual` around a tracked value would be cut off its tape, so there is no common type +Base.promote_rule(::Type{TrackedReal{V1,D,O}}, ::Type{Dual{T,V2,N}}) where {V1,D,O,T,V2<:TrackedReal,N} = Union{} +Base.promote_rule(::Type{Dual{T,V1,N}}, ::Type{TrackedReal{V2,D,O}}) where {T,V1<:TrackedReal,N,V2,D,O} = Union{} Base.promote_rule(::Type{TrackedReal{V1,D1,O1}}, ::Type{TrackedReal{V2,D2,O2}}) where {V1,V2,D1,D2,O1,O2} = TrackedReal{promote_type(V1,V2),promote_type(D1,D2),Nothing} diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 0b54b682..937a16b6 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -77,6 +77,18 @@ end @test ReverseDiff.gradient(x -> sum(x .* d), a) ≈ Matrix(d) @test ReverseDiff.gradient(x -> sum(x .* u), a) ≈ Matrix(u) @test ReverseDiff.gradient(x -> sum(x .* s), v) ≈ s + + # a closure takes the scalar rules, and still broadcasts a static value statically + c = 3.0 + tp = InstructionTape() + xs = track(SVector(1.0, 2.0), tp) + y, z = (t -> t * c).(xs), ((t, u) -> t * u * c).(xs, SVector(1.0, 2.0)) + @test y isa SVector + @test z isa SVector + ReverseDiff.seed!(sum(y) + sum(z)) + ReverseDiff.reverse_pass!(tp) + @test deriv(xs) == [6.0, 9.0] + @test ReverseDiff.gradient(x -> sum((t -> t * c).(x)), SVector(1.0, 2.0)) == [3.0, 3.0] end @testset "no `BroadcastStyle` ambiguities" begin @@ -130,6 +142,11 @@ end a = rand(3) @test ReverseDiff.gradient(f, a) ≈ exp.(a) @test ReverseDiff.hessian(f, a) ≈ zeros(3, 3) + + # arrays too + s = ReverseDiff.SkipOptimize(v -> sum(exp, v)) + @test ReverseDiff.gradient(v -> sum(v) * s(v), a) ≈ fill(sum(exp, a), 3) + @test ReverseDiff.hessian(v -> sum(v) * s(v), a) == zeros(3, 3) end @testset "a type broadcast as a function does not force the fallback path" begin @@ -238,6 +255,16 @@ end @test y isa TrackedArray @test value(y) == [0.0, 2.0] @test ReverseDiff.gradient(x -> sum(relu.(x)), a) == [0.0, 1.0] + + # recorded with every element on the constant branch, replayed on the other + relu0(t) = t > 0 ? t : 0.0 + for f in (x -> sum(relu.(x)) + sum(x), x -> sum(relu0.(x)) + sum(x), + x -> sum(ifelse.(x .> 0, x, 0.0)) + sum(x)) + tape = ReverseDiff.GradientTape(f, [-1.0, -2.0]) + @test ReverseDiff.gradient!(tape, [1.0, 2.0]) == [2.0, 2.0] + @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [1.0, 2.0]) == [2.0, 2.0] + end + @test relu0.(track([-1.0, -2.0], InstructionTape())) isa TrackedArray end @testset "a non-`Real` result keeps its partials" begin @@ -262,15 +289,13 @@ end # `Tag`s are ordered by first use, so pin the order before relying on it @test ForwardDiff.:≺(tagA, tagB) - concrete = ForwardDiff.Dual{tagA,Float64,1}[ForwardDiff.Dual{tagA}(2.0, 3.0)] - widened = Union{Float64,ForwardDiff.Dual{tagA,Float64,1}}[ - 1.0, ForwardDiff.Dual{tagA}(2.0, 3.0)] - nested = Union{Float64,ForwardDiff.Dual{tagB,Float64,1}}[1.0] + # `trackresults` asks for the element type `df` infers on `vals` + outer = [ForwardDiff.Dual{tagA}(2.0, 3.0)] + nested = [ForwardDiff.Dual{tagB}(2.0, 3.0)] - @test ReverseDiff.trackresults(tagB, concrete, identity, identity, (), ()) === concrete - @test ReverseDiff.trackresults(tagB, widened, identity, identity, (), ()) === widened + @test ReverseDiff.trackresults(tagB, outer, i -> outer[i], identity, (), ([1],)) === outer @test_throws ForwardDiff.DualMismatchError ReverseDiff.trackresults( - tagA, nested, identity, identity, (), ()) + tagA, nested, i -> nested[i], identity, (), ([1],)) @test ReverseDiff.getpartial(tagA, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 2.0 @test ReverseDiff.getpartial(tagA, 1.0, 1) == 0.0 @@ -287,16 +312,51 @@ end @test_throws ArgumentError(msg) ForwardDiff.derivative(g, 3.0) - # a widened element type is checked member by member + # an abstract element type is checked element by element g2(a) = sum(ReverseDiff.gradient(x -> sum(x .* Union{Float64,typeof(a)}[1.0, a]), [1.0, 2.0])) E2 = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g2),Float64},Float64,1} msg2 = "a broadcast argument with element type $E2 carries a perturbation that a derivative of type Float64 cannot hold" @test_throws ArgumentError(msg2) ForwardDiff.derivative(g2, 3.0) + g3(a) = sum(ReverseDiff.gradient(x -> sum(ifelse.(x .> 5, x, Real[1.0, a])), [1.0, 2.0])) + E3 = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g3),Float64},Float64,1} + msg3 = "a broadcast argument with element type $E3 carries a perturbation that a derivative of type Float64 cannot hold" + @test_throws ArgumentError(msg3) ForwardDiff.derivative(g3, 3.0) # a tape whose derivatives are themselves `Dual`s can hold it, so it is left alone h(a) = ReverseDiff.gradient(x -> sum(x .* a), [ForwardDiff.Dual(1.0, 0.0)]) @test h(3.0) == [ForwardDiff.Dual(3.0, 0.0)] + # also when the tag without the perturbation was created first + ReverseDiff.gradient(x -> sum(atan.(x, 2.0)), [1.0, 2.0]) + k(a) = sum(ReverseDiff.gradient(x -> sum(atan.(x, a)), [1.0, 2.0], + ReverseDiff.GradientConfig([1.0, 2.0], typeof(a)))) + xs = [1.0, 2.0] + @test ForwardDiff.derivative(k, 3.0) ≈ sum(@. (xs^2 - 9) / (xs^2 + 9)^2) + + # the scalar rules bury the tracked value in the `Dual`, which must not give a zero derivative + msg4 = "ForwardDiff cannot differentiate through ReverseDiff (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)" + for g4 in (a -> sum(ReverseDiff.gradient(x -> x[1] * a, [1.0, 2.0])), + a -> sum(ReverseDiff.gradient(x -> sum(x .* Real[1.0, a]), [1.0, 2.0])), + a -> sum(ReverseDiff.jacobian(x -> [x[1] * a, x[2]], [1.0, 2.0]))) + @test_throws ArgumentError(msg4) ForwardDiff.derivative(g4, 3.0) + end + # storing the `Dual` as a tracked number would cut it off the tape + msg5 = "this nesting of ForwardDiff and ReverseDiff is not supported: a `Dual` of tracked numbers cannot be converted to a tracked number (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)" + function g5(a) + return sum(ReverseDiff.gradient([1.0, 2.0]) do x + v = [x[1], x[2]] + v[1] = x[1] * a + return sum(v) + end) + end + @test_throws ArgumentError(msg5) ForwardDiff.derivative(g5, 3.0) + + # the other way around is fine, also next to tracked numbers + @test ReverseDiff.gradient(x -> ForwardDiff.derivative(a -> x[1] * a^2, 2.0), [1.0, 2.0]) == + [4.0, 0.0] + @test ReverseDiff.gradient(x -> ForwardDiff.derivative(a -> sum([x[1] * a^2, x[2]]), 2.0), + [1.0, 2.0]) == [4.0, 0.0] + # also when the perturbation comes from an enclosing broadcast f(x) = sum(broadcast(a -> ReverseDiff.gradient(y -> sum(y .* a), [1.0])[1], x)) @test_throws ArgumentError ReverseDiff.gradient(f, [2.0, 3.0]) From e4a76282d6f3910b8455dd73eb7c5bedfbebc3aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 29 Sep 2026 10:34:21 +0200 Subject: [PATCH 23/36] Capture the broadcast tag by instance so Julia 1.10 infers the dualized call Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index ea6e885f..b860f1ac 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -228,19 +228,19 @@ skipvalue(x::Dual{T}) where {T<:ForwardDiff.Tag{<:BroadcastTag}} = skipvalue(For _, vals, _ = splitargs(argvals) slots, positions, valP = trackedslots(targs) # keyed on every argument's type, so an enclosing differentiation's `Dual` makes a newer tag - T = typeof(ForwardDiff.Tag(BroadcastTag{F}(), typeof(argvals))) + tag = ForwardDiff.Tag(BroadcastTag{F}(), typeof(argvals)) # `broadcast` calls `df` elementwise, so it receives one scalar per argument function df(x::Vararg{Any,N}) where {N} - dx = map((slot, xi) -> dualize(T, slot, valP, xi), slots, x) + dx = map((slot, xi) -> dualize(typeof(tag), slot, valP, xi), slots, x) return splatcall(f, dx, untracked, inds) end # known partials leave nothing to read off a `Dual`, so `f` is evaluated undualized vf(x::Vararg{Any,N}) where {N} = splatcall(f, x, untracked, inds) entries = knownpartials(f, args, positions) if entries === nothing - return trackresults(T, broadcast(df, vals...), df, vf, targs, vals) + return trackresults(typeof(tag), broadcast(df, vals...), df, vf, targs, vals) else - return recordresults(T, KnownPartials(entries), df, vf, targs, vals) + return recordresults(typeof(tag), KnownPartials(entries), df, vf, targs, vals) end end From a2642aeeef10abe20cf1f95b3d679eba251b9250 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 29 Sep 2026 10:52:58 +0200 Subject: [PATCH 24/36] Keep static broadcast results writable, simplify `_mayhidetracked`, and check abstract argument elements without dispatch - store the results of a static array via `similar`, which gives a `SizedArray` for a non-isbits element type where `convert` gives an `MArray` that rejects `setindex!` - replace the `@generated` `try`/`catch` around `fieldcount` with `isconcretetype` - branch on `isa` instead of dispatching on every element of an abstractly typed argument - stop pinning `MethodError`s of the internal `seed!` - test that the cached partials have a concrete element type Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 26 ++++++++++++++------------ test/TrackedTests.jl | 25 ------------------------- test/derivatives/BroadcastTests.jl | 11 +++++++++++ 3 files changed, 25 insertions(+), 37 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index b860f1ac..3cd477d7 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -31,16 +31,8 @@ mayhidetracked(b::Broadcasted) = mayhidetracked(b.f) || any(mayhidetracked, b.ar # into what a container holds rather than ask about the container _mayhidetracked(::Type{<:NotTracked}) = false _mayhidetracked(::Type{<:AbstractArray{F}}) where {F} = _mayhidetracked(F) -@generated function _mayhidetracked(::Type{F}) where {F} - # `fieldcount` errors for types without a definite number of fields, such as - # `Type{T}` and abstract types; be conservative in that case. - hasfields = try - fieldcount(F) > 0 - catch - true - end - return :($hasfields) -end +# a type that is not concrete, such as `Type{T}` or an abstract type, may have fields +_mayhidetracked(::Type{F}) where {F} = !isconcretetype(F) || fieldcount(F) > 0 struct TrackedStyle{N} <: AbstractArrayStyle{N} end @@ -265,7 +257,13 @@ function checkargvalues(::Type{D}, v) where {D} if isconcretetype(eltype(v)) checkargtags(D, eltype(v)) else - foreach(x -> checkargvalues(D, x), v) + for x in v + if x isa Dual + checkargtags(D, typeof(x)) + elseif !(x isa Real) + checkargvalues(D, x) + end + end end return nothing end @@ -306,10 +304,14 @@ end checktags(T, E) return results else - return recordresults(T, convert(AbstractArray{E}, results), df, vf, targs, vals) + return recordresults(T, writable(results, E), df, vf, targs, vals) end end +# the replay writes into the results, which a static array may reject +writable(results::AbstractArray, ::Type{E}) where {E} = convert(AbstractArray{E}, results) +writable(results::StaticArray, ::Type{E}) where {E} = copyto!(similar(results, E), results) + @noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) input = instruction.input output = instruction.output diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index 050795dc..682a73ef 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -389,16 +389,11 @@ ta = TrackedArray(varr, darr, InstructionTape()) tr = TrackedReal(v, d) trs = Any[tr, ta[1], ta[2]] -@test_throws MethodError ReverseDiff.seed!(nothing) -@test_throws MethodError ReverseDiff.seed!([1,2,3]) @test ReverseDiff.seed!(1.0) === nothing @test ReverseDiff.unseed!(nothing) === nothing @test ReverseDiff.unseed!([1,2,3]) === nothing @test ReverseDiff.unseed!(1.0) === nothing -@test_throws MethodError ReverseDiff.seed!(ta) -@test ReverseDiff.value(ta) === varr == varr_copy -@test ReverseDiff.deriv(ta) === darr == darr_copy ReverseDiff.unseed!(ta) @test ReverseDiff.value(ta) === varr == varr_copy @test ReverseDiff.deriv(ta) === darr == fill!(similar(darr), 0) @@ -416,16 +411,6 @@ ReverseDiff.deriv!(tr, d) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === d -@test_throws MethodError ReverseDiff.seed!(trs) -@test ReverseDiff.value(tr) === v -@test ReverseDiff.deriv(tr) === d -@test ReverseDiff.value(ta) === varr == varr_copy -@test ReverseDiff.deriv(ta) === darr == darr_copy -@test ReverseDiff.value(trs[2]) === varr[1] -@test ReverseDiff.deriv(trs[2]) === darr[1] -@test ReverseDiff.value(trs[3]) === varr[2] -@test ReverseDiff.deriv(trs[3]) === darr[2] - ReverseDiff.unseed!(trs) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === 0.0 @@ -445,16 +430,6 @@ ReverseDiff.pull_deriv!(trs) @test ReverseDiff.deriv(trs[2]) === darr[1] @test ReverseDiff.deriv(trs[3]) === darr[2] -@test_throws MethodError ReverseDiff.seed!((tr, ta, trs)) -@test ReverseDiff.value(tr) === v -@test ReverseDiff.deriv(tr) === d -@test ReverseDiff.value(ta) === varr == varr_copy -@test ReverseDiff.deriv(ta) === darr == darr_copy -@test ReverseDiff.value(trs[2]) === varr[1] -@test ReverseDiff.deriv(trs[2]) === darr[1] -@test ReverseDiff.value(trs[3]) === varr[2] -@test ReverseDiff.deriv(trs[3]) === darr[2] - ReverseDiff.unseed!((tr, ta, trs)) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === 0.0 diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 937a16b6..2a12c6d6 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -265,6 +265,17 @@ end @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [1.0, 2.0]) == [2.0, 2.0] end @test relu0.(track([-1.0, -2.0], InstructionTape())) isa TrackedArray + + # a static array whose widened element type is not isbits + tape = ReverseDiff.GradientTape(x -> sum(relu.(x)), MVector(-1.0, 2.0)) + @test ReverseDiff.gradient!(tape, MVector(1.0, 2.0)) == [1.0, 1.0] +end + +@testset "the cached partials have a concrete element type" begin + # the inferred element type decides what is recorded + tp = InstructionTape() + track(rand(3), tp) .^ 2 + @test isconcretetype(eltype(first(tp[1].cache))) end @testset "a non-`Real` result keeps its partials" begin From 171b7c22c83785b7285d04cafe50cf48e688ea3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 2 Oct 2026 14:24:04 +0200 Subject: [PATCH 25/36] Reject 0-dimensional `TrackedArray` destinations, test tags end to end, and finish tapes before reading them - `copyto!` into a 0-dimensional `TrackedArray` was ambiguous with `Base`'s 0-dimensional method instead of throwing the destination error. - Replace the tests of `trackresults` and `getpartial` with nestings of ForwardDiff inside and around the broadcast. - The broadcast tests finish their tapes before reading or replaying them (#311). - Name the replayed function `replayf` and clarify two comments. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 15 +++--- test/derivatives/BroadcastTests.jl | 77 ++++++++++++++++-------------- 2 files changed, 51 insertions(+), 41 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 4017eb70..c0f17187 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -27,8 +27,7 @@ mayhidetracked(b::ForwardOptimize) = mayhidetracked(b.f) mayhidetracked(b::SkipOptimize) = mayhidetracked(b.f) mayhidetracked(b::Broadcasted) = mayhidetracked(b.f) || any(mayhidetracked, b.args) -# below the argument nothing is seeded, so a tracked value is hidden wherever it sits: recurse -# into what a container holds rather than ask about the container +# nothing inside an argument is seeded, so ask about the element type, not the container _mayhidetracked(::Type{<:NotTracked}) = false _mayhidetracked(::Type{<:AbstractArray{F}}) where {F} = _mayhidetracked(F) # a type that is not concrete, such as `Type{T}` or an abstract type, may have fields @@ -39,6 +38,7 @@ struct TrackedStyle{N} <: AbstractArrayStyle{N} end (::Type{<:TrackedStyle})(::Val{N}) where {N} = TrackedStyle{N}() Broadcast.BroadcastStyle(::Type{<:TrackedReal}) = TrackedStyle{0}() +# only an element type that fixes `D`, since `getouttype` reads it off Broadcast.BroadcastStyle(::Type{<:AbstractArray{<:TrackedReal{<:Any, D}, N}}) where {D, N} = TrackedStyle{N}() @@ -101,6 +101,9 @@ _no_tracked_dest() = throw(ArgumentError("`TrackedArray`s do not support `setind Base.copyto!(::TrackedArray, ::Broadcasted{<:TrackedStyle}) = _no_tracked_dest() Base.copyto!(::TrackedArray, ::Broadcasted{<:DefaultArrayStyle}) = _no_tracked_dest() +# more specific than `Base`'s 0-dimensional method +Base.copyto!(::TrackedArray, ::Broadcasted{TrackedStyle{0}}) = _no_tracked_dest() +Base.copyto!(::TrackedArray, ::Broadcasted{DefaultArrayStyle{0}}) = _no_tracked_dest() getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D @@ -277,12 +280,12 @@ end @inline function recordresults(::Type{T}, results, df, vf, targs, vals) where {T} D = mapreduce(getouttype, promote_type, targs) foreach(v -> checkargvalues(D, v), vals) - g, outvalue = replaycache(T, results, df, vf, vals) + replayf, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) out = track(outvalue, D, tp) _, positions, n = trackedslots(targs) bounds = map(p -> index_bound(getat(targs, p), out), positions) - cache = (results, g, T(), map(tuple, positions, ntuple(Val, n), bounds)) + cache = (results, replayf, T(), map(tuple, positions, ntuple(Val, n), bounds)) record!(tp, SpecialInstruction, ∇broadcast, targs, out, cache) return out end @@ -365,9 +368,9 @@ _increment_deriv!(::Type, x, out_deriv, p::PartialsWithArgs, ::Val{k}, bound...) @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) input, output = instruction.input, instruction.output - results, df, tag, _ = instruction.cache + results, replayf, tag, _ = instruction.cache foreach(pull_value!, input) - _replay!(typeof(tag), value(output), results, df, map(value, input)) + _replay!(typeof(tag), value(output), results, replayf, map(value, input)) return nothing end diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 84fd6d5f..d56c3771 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -28,6 +28,9 @@ Base.getindex(f::Foreign0d, ::CartesianIndex{0}) = f.x Base.Broadcast.BroadcastStyle(::Type{<:Foreign0d}) = ForeignStyle() Base.similar(::Broadcasted{ForeignStyle}, ::Type{T}, axs) where {T} = similar(Array{T}, axs) +# at top level, since a local function that captures itself takes the scalar rules +inner(t) = ForwardDiff.Dual{typeof(ForwardDiff.Tag(inner, typeof(t)))}(t, one(t)) + ############################################################################################ @testset "`BroadcastStyle` tracks dimensionality" begin @@ -88,7 +91,7 @@ end @test y isa SVector @test z isa SVector ReverseDiff.seed!(sum(y) + sum(z)) - ReverseDiff.reverse_pass!(tp) + ReverseDiff.reverse_pass!(ReverseDiff.finish!(tp)) @test deriv(xs) == [6.0, 9.0] @test ReverseDiff.gradient(x -> sum((t -> t * c).(x)), SVector(1.0, 2.0)) == [3.0, 3.0] end @@ -121,10 +124,11 @@ end ) tp1, tp2 = InstructionTape(), InstructionTape() x, y = dotted(tp1), called(tp2) + instrs1, instrs2 = take_recorded!(tp1), take_recorded!(tp2) @test value(x) == value(y) - @test length(tp1) == length(tp2) == 1 - @test tp1[1].func === tp2[1].func + @test length(instrs1) == length(instrs2) == 1 + @test instrs1[1].func === instrs2[1].func end end @@ -133,7 +137,7 @@ end x = track(rand(3, 3), tp) y = broadcast(ReverseDiff.ForwardOptimize(exp), x) @test y isa TrackedArray - @test length(tp) == 1 + @test length(take_recorded!(tp)) == 1 # `@skip` results are untracked, as with `map` and scalars tp = InstructionTape() @@ -141,7 +145,7 @@ end y = broadcast(ReverseDiff.SkipOptimize(exp), x) @test y isa Matrix{Float64} @test y == exp.(value(x)) - @test isempty(tp) + @test isempty(take_recorded!(tp)) # fused, and at every order f(v) = sum(v .* ReverseDiff.@skip(exp).(v)) @@ -160,10 +164,11 @@ end x = track(rand(3), tp) y = Real.(x) + instrs = take_recorded!(tp) @test y isa TrackedArray - @test length(tp) == 1 - @test tp[1].func === ReverseDiff.∇broadcast + @test length(instrs) == 1 + @test instrs[1].func === ReverseDiff.∇broadcast @test ReverseDiff.gradient(v -> sum(Real.(v)), [1.0, 2.0]) == [1.0, 1.0] end @@ -182,7 +187,7 @@ end sum((x .+ y) .* sin.(x)) # one instruction for the fused expression, one for `sum` - @test length(tp) == 2 + @test length(take_recorded!(tp)) == 2 end @testset "`TrackedArray`s are rejected as broadcast destinations" begin @@ -193,8 +198,9 @@ end x = track(copy(a), tp) @test_throws ArgumentError(msg) track(zeros(4), tp) .= exp.(x) @test_throws ArgumentError(msg) track(zeros(4), tp) .= a + @test_throws ArgumentError(msg) track(fill(1.0), tp) .= 2.0 # nothing may be recorded before the failure - @test isempty(tp) + @test isempty(take_recorded!(tp)) # an untracked container of tracked elements is a valid destination dest = Vector{TrackedReal{Float64, Float64, Nothing}}(undef, 4) @@ -220,13 +226,16 @@ end x = map(xi -> track(xi, tp), a) y = x .* b + s = sum(y) + ReverseDiff.finish!(tp) @test y isa TrackedArray @test value(y) ≈ a .* b - @test length(tp) == 1 - @test tp[1].func === ReverseDiff.∇broadcast + # one instruction for the broadcast, one for `sum` + @test length(tp) == 2 + @test ReverseDiff.instructions(tp)[1].func === ReverseDiff.∇broadcast - ReverseDiff.seed!(sum(y)) + ReverseDiff.seed!(s) ReverseDiff.reverse_pass!(tp) # `deriv.(x)` would be traced like any other broadcast, as it is for a `TrackedArray` @test map(deriv, x) ≈ b @@ -283,7 +292,7 @@ end # the inferred element type decides what is recorded tp = InstructionTape() track(rand(3), tp) .^ 2 - @test isconcretetype(eltype(first(tp[1].cache))) + @test isconcretetype(eltype(first(only(take_recorded!(tp)).cache))) end @testset "a non-`Real` result keeps its partials" begin @@ -303,26 +312,23 @@ end end @testset "a foreign tag is never read as our own" begin - tagA = typeof(ForwardDiff.Tag(sin, Float64)) - tagB = typeof(ForwardDiff.Tag(cos, Float64)) - # `Tag`s are ordered by first use, so pin the order before relying on it - @test ForwardDiff.:≺(tagA, tagB) - - # `trackresults` asks for the element type `df` infers on `vals` - outer = [ForwardDiff.Dual{tagA}(2.0, 3.0)] - nested = [ForwardDiff.Dual{tagB}(2.0, 3.0)] - - @test ReverseDiff.trackresults(tagB, outer, i -> outer[i], identity, (), ([1],)) === outer - @test_throws ForwardDiff.DualMismatchError ReverseDiff.trackresults( - tagA, nested, i -> nested[i], identity, (), ([1],) + # a `Dual` built inside the function buries the broadcast's own partial + @test_throws ForwardDiff.DualMismatchError ReverseDiff.gradient( + x -> sum(inner.(x)), [1.0, 2.0] + ) + # also on one branch only, which is recorded + @test_throws ForwardDiff.DualMismatchError ReverseDiff.gradient( + x -> sum((t -> t > 1.5 ? inner(t) : t).(x)), [1.0, 2.0] ) - @test ReverseDiff.getpartial(tagA, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 2.0 - @test ReverseDiff.getpartial(tagA, 1.0, 1) == 0.0 - @test ReverseDiff.getpartial(tagB, ForwardDiff.Dual{tagA}(1.0, 2.0), 1) == 0.0 - @test_throws ForwardDiff.DualMismatchError ReverseDiff.getpartial( - tagA, ForwardDiff.Dual{tagB}(1.0, 2.0), 1 + # an enclosing differentiation's `Dual` is constant in our arguments + xs = [1.0, 3.0] + m(a) = sum( + ReverseDiff.gradient( + x -> sum(ifelse.(x .> 2, x, a)), xs, ReverseDiff.GradientConfig(xs, typeof(a)) + ) ) + @test ForwardDiff.derivative(m, 1.0) == 0.0 end @testset "a perturbation the derivative cannot hold is rejected (#67, #168)" begin @@ -462,7 +468,7 @@ end actual = map(value, wrap(track(d, tp), track(e, tp), track(A, tp))) @test actual == expected - @test isempty(tp) + @test isempty(take_recorded!(tp)) # the test covers every `StructuredMatrix` type covered = map(nameof ∘ typeof, expected) @@ -476,9 +482,10 @@ end tp = InstructionTape() y = track(copy(a), tp) .* c + instrs = take_recorded!(tp) @test y isa TrackedArray - @test length(tp) == 1 - @test tp[1].func === ReverseDiff.∇broadcast + @test length(instrs) == 1 + @test instrs[1].func === ReverseDiff.∇broadcast @test ReverseDiff.gradient(x -> sum(x .* c), a) == collect(c) tape = ReverseDiff.GradientTape(x -> sum(exp.(x .* c)), a) @@ -494,7 +501,7 @@ end tp = InstructionTape() track(copy(a), tp) .* first.(((1.0, 2.0), (3.0, 4.0), (5.0, 6.0))) - @test length(tp) == 1 + @test length(take_recorded!(tp)) == 1 end @testset "a `TrackedReal` without a tape is a constant" begin @@ -505,7 +512,7 @@ end tp = InstructionTape() x = track(copy(a), tp) ReverseDiff.seed!(sum(f(x))) - ReverseDiff.reverse_pass!(tp) + ReverseDiff.reverse_pass!(ReverseDiff.finish!(tp)) @test deriv(x) ≈ expected end end From 9fda02ab9e979dabf82908c05f358429718fb2d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 2 Oct 2026 17:33:17 +0200 Subject: [PATCH 26/36] Cover the remaining broadcast paths and tidy up the review leftovers - Recurse into argument values only through tuples and arrays, so a non-concrete tuple holding `nothing` no longer throws a `MethodError` - Capture the broadcast tag rather than its type in the reverse pass - Drop the `copyto!` methods for `DefaultArrayStyle`, which `.=` never reaches - Restore master's `return nothing` in `reverse_mul!` and print `TrackedReal`s without building a string first - Test `NotTracked` scalars and arguments, the `Union` branch of `checktags`, and `map` over skipped functions; use `gradient` for the tapeless `TrackedReal` test Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 8 +++----- src/derivatives/linalg/arithmetic.jl | 3 ++- src/tracked.jl | 18 +++++++++++++---- test/derivatives/BroadcastTests.jl | 30 ++++++++++++++++++++++------ test/derivatives/ElementwiseTests.jl | 19 ++++++++++++++++++ 5 files changed, 62 insertions(+), 16 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index c0f17187..e2c1695d 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -99,11 +99,10 @@ end _no_tracked_dest() = throw(ArgumentError("`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead.")) +# `.=` combines the destination's style, so a `TrackedArray` destination is always `TrackedStyle` Base.copyto!(::TrackedArray, ::Broadcasted{<:TrackedStyle}) = _no_tracked_dest() -Base.copyto!(::TrackedArray, ::Broadcasted{<:DefaultArrayStyle}) = _no_tracked_dest() # more specific than `Base`'s 0-dimensional method Base.copyto!(::TrackedArray, ::Broadcasted{TrackedStyle{0}}) = _no_tracked_dest() -Base.copyto!(::TrackedArray, ::Broadcasted{DefaultArrayStyle{0}}) = _no_tracked_dest() getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D @@ -269,7 +268,7 @@ function checkargvalues(::Type{D}, v) where {D} for x in v if x isa Dual checkargtags(D, typeof(x)) - elseif !(x isa Real) + elseif x isa Union{Tuple, AbstractArray} checkargvalues(D, x) end end @@ -322,11 +321,10 @@ end output = instruction.output output_deriv = deriv(output) results, _, tag, targets = instruction.cache - T = typeof(tag) partials = reversepartials(results, input) foreach(targets) do (p, k, bound) x = getat(input, p) - istracked(x) && _br_add_to_deriv!(T, x, k, output_deriv, partials, bound) + istracked(x) && _br_add_to_deriv!(typeof(tag), x, k, output_deriv, partials, bound) end unseed!(output) return nothing diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index 1459ef57..1a898efa 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -286,7 +286,8 @@ function reverse_mul!(output, output_deriv, a, b, a_tmp, b_tmp) increment_deriv!(a, structured_mul!(a_tmp, output_deriv, transpose(value(b)))) end end - return istracked(b) && increment_deriv!(b, structured_mul!(b_tmp, transpose(value(a)), output_deriv)) + istracked(b) && increment_deriv!(b, structured_mul!(b_tmp, transpose(value(a)), output_deriv)) + return nothing end ## zero diff --git a/src/tracked.jl b/src/tracked.jl index 65b190f5..4487a4aa 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -578,8 +578,18 @@ end idstr(x) = string(objectid(x), base = 62)[1:3] function Base.show(io::IO, t::TrackedReal) - tape_id = hastape(t) ? idstr(t.tape) : "---" - origin_id = hasorigin(t) ? "$(t.index), $(idstr(t.origin))" : "---" - id = idstr(t) - return print(io, "TrackedReal<$(id)>($(value(t)), $(deriv(t)), $(tape_id), $(origin_id))") + print(io, "TrackedReal<", idstr(t), ">(", value(t), ", ", deriv(t), ", ") + if hastape(t) + print(io, idstr(t.tape)) + else + print(io, "---") + end + print(io, ", ") + if hasorigin(t) + print(io, t.index, ", ", idstr(t.origin)) + else + print(io, "---") + end + print(io, ")") + return nothing end diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index d56c3771..9c18a784 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -31,6 +31,11 @@ Base.similar(::Broadcasted{ForeignStyle}, ::Type{T}, axs) where {T} = similar(Ar # at top level, since a local function that captures itself takes the scalar rules inner(t) = ForwardDiff.Dual{typeof(ForwardDiff.Tag(inner, typeof(t)))}(t, one(t)) +# a struct with fields, which `NotTracked` declares constant +struct Scale + s::Float64 +end + ############################################################################################ @testset "`BroadcastStyle` tracks dimensionality" begin @@ -320,6 +325,10 @@ end @test_throws ForwardDiff.DualMismatchError ReverseDiff.gradient( x -> sum((t -> t > 1.5 ? inner(t) : t).(x)), [1.0, 2.0] ) + # and next to a constant branch, which is not + @test_throws ForwardDiff.DualMismatchError ReverseDiff.gradient( + x -> sum((t -> t > 1.5 ? inner(t) : 1.0).(x)), [1.0, 2.0] + ) # an enclosing differentiation's `Dual` is constant in our arguments xs = [1.0, 3.0] @@ -502,6 +511,10 @@ end tp = InstructionTape() track(copy(a), tp) .* first.(((1.0, 2.0), (3.0, 4.0), (5.0, 6.0))) @test length(take_recorded!(tp)) == 1 + + # an abstract element type is checked element by element, skipping non-numbers + g = x -> sum(x .* first.(((1.0, nothing), (2, nothing), (3.0, nothing)))) + @test ReverseDiff.gradient(g, a) == [1.0, 2.0, 3.0] end @testset "a `TrackedReal` without a tape is a constant" begin @@ -509,18 +522,23 @@ end c = TrackedReal(2.0, 0.0) for (f, expected) in ((x -> x .* c, [2.0, 2.0]), (x -> exp.(x .* c), 2 .* exp.(2 .* a))) - tp = InstructionTape() - x = track(copy(a), tp) - ReverseDiff.seed!(sum(f(x))) - ReverseDiff.reverse_pass!(ReverseDiff.finish!(tp)) - @test deriv(x) ≈ expected + @test ReverseDiff.gradient(x -> sum(f(x)), a) ≈ expected end end @testset "`NotTracked`" begin f = ReverseDiff.NotTracked(t -> 2t) + a = [1.0, 2.0, 3.0] @test ReverseDiff.gradient(x -> sum(f.(x)), rand(4)) == fill(2.0, 4) + # a scalar broadcast takes the scalar rules, calling the wrapper itself + @test ReverseDiff.gradient(x -> f.(x[1]), a) == [2.0, 0.0, 0.0] + + # as an argument, a struct is broadcast as a scalar and an array elementwise + g = (t, s) -> t * s.s + @test ReverseDiff.gradient(x -> sum(g.(x, ReverseDiff.NotTracked(Scale(2.0)))), a) == + fill(2.0, 3) + @test ReverseDiff.gradient(x -> sum(g.(x, ReverseDiff.NotTracked(Scale.(a)))), a) == a end @testset "a partial known in closed form stores nothing" begin @@ -578,7 +596,7 @@ end @test ReverseDiff.gradient(t -> sum(t' .* t'), v) == 2 .* v # `-x/y^2` is no argument of the broadcast, so it is read off a `Dual` - @test ReverseDiff.gradient(t -> sum(w ./ t), v) == -w ./ v .^ 2 + @test ReverseDiff.gradient(t -> sum(w ./ t), v) ≈ -w ./ v .^ 2 # an outer broadcast still clamps the reverse pass to the argument's own shape @test ReverseDiff.gradient(t -> sum(t .+ w'), v) == fill(3.0, 3) diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index 822da3ff..da082efe 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -528,4 +528,23 @@ end end end +@testset "`map` with a skipped function is untracked" begin + a, b, c = [1.0, 2.0, 3.5], [2.0, 2.0, 1.0], 2.0 + + tp = InstructionTape() + at, bt = track(a, tp), track(b, tp) + + for fsym in ReverseDiff.SKIPPED_UNARY_SCALAR_FUNCS + f = eval(fsym) + @test map(f, at) == map(f, a) + end + for fsym in ReverseDiff.SKIPPED_BINARY_SCALAR_FUNCS + f = eval(fsym) + for (x, y) in ((at, bt), (a, bt), (at, b), (c, bt), (at, c)) + @test map(f, x, y) == map(f, ReverseDiff.value(x), ReverseDiff.value(y)) + end + end + @test isempty(take_recorded!(tp)) +end + end # module From d27fbef48c6a2710c0b6c05fd2dc37b3f337eb81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sat, 3 Oct 2026 00:34:08 +0200 Subject: [PATCH 27/36] Catch `TrackedArray` broadcast destinations by type, and test the static fallback for `StaticArray` `.=` hands `copyto!` the source's style, so an untracked source never reached the `TrackedStyle` methods. A tracked static value is stored as an `MVector`, which the fallback keeps. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 6 ++---- test/derivatives/BroadcastTests.jl | 6 ++++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index e2c1695d..e16d23d6 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -99,10 +99,8 @@ end _no_tracked_dest() = throw(ArgumentError("`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead.")) -# `.=` combines the destination's style, so a `TrackedArray` destination is always `TrackedStyle` -Base.copyto!(::TrackedArray, ::Broadcasted{<:TrackedStyle}) = _no_tracked_dest() -# more specific than `Base`'s 0-dimensional method -Base.copyto!(::TrackedArray, ::Broadcasted{TrackedStyle{0}}) = _no_tracked_dest() +# `.=` hands `copyto!` the source's style, not the destination's, so catch the destination here +Broadcast.materialize!(::TrackedArray, ::Broadcasted) = _no_tracked_dest() getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 9c18a784..19bb28c7 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -93,8 +93,8 @@ end tp = InstructionTape() xs = track(SVector(1.0, 2.0), tp) y, z = (t -> t * c).(xs), ((t, u) -> t * u * c).(xs, SVector(1.0, 2.0)) - @test y isa SVector - @test z isa SVector + @test y isa StaticArray + @test z isa StaticArray ReverseDiff.seed!(sum(y) + sum(z)) ReverseDiff.reverse_pass!(ReverseDiff.finish!(tp)) @test deriv(xs) == [6.0, 9.0] @@ -204,6 +204,8 @@ end @test_throws ArgumentError(msg) track(zeros(4), tp) .= exp.(x) @test_throws ArgumentError(msg) track(zeros(4), tp) .= a @test_throws ArgumentError(msg) track(fill(1.0), tp) .= 2.0 + # a foreign style brings its own `copyto!` + @test_throws ArgumentError(msg) track(zeros(2), tp) .= SVector(1.0, 2.0) # nothing may be recorded before the failure @test isempty(take_recorded!(tp)) From 986244306b4b1eab55bd4326351079976a147bdf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sun, 4 Oct 2026 01:49:37 +0200 Subject: [PATCH 28/36] Keep widened broadcast values on replay, drop the dead `Union` branch, and test `show` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `map` narrowed a widened broadcast's value buffer to the recorded values' type, so replaying `t > 0 ? t : 0` at non-integer inputs threw an `InexactError`. `combine_eltypes` widens to `Real`, never to a `Union`, so `checktags` loses its unreachable branch. The `nothing`-in-tuple test now reaches `∇broadcast` instead of the scalar fallback. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 27 +++++++++++++-------------- test/TrackedTests.jl | 8 ++++++++ test/derivatives/BroadcastTests.jl | 13 +++++++++---- 3 files changed, 30 insertions(+), 18 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index e16d23d6..0e379320 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -77,8 +77,8 @@ function remove_not_tracked(b::Broadcasted{style}) where {style} end function Base.copy(_bc::Broadcasted{<:TrackedStyle}) - # scalars take the scalar derivative rules; ask the axes, since `TrackedStyle{Any}` carries - # no dimension and `LinearAlgebra` may pass an uninstantiated `Broadcasted` + # scalars take the scalar derivative rules. Ask the axes, since `TrackedStyle{Any}` carries + # no dimension and `LinearAlgebra` may pass an uninstantiated `Broadcasted`. if axes(_bc) isa Tuple{} return _bc[CartesianIndex()] end @@ -238,8 +238,10 @@ end # the cache carries what the replay needs: `df` to recompute the stored partials, or, where # they are known already, `vf` for the values alone -replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} = - (df, map(y -> ForwardDiff.value(T, y), results)) +function replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} + outvalue = similar(results, ForwardDiff.valtype(T, eltype(results))) + return df, map!(y -> ForwardDiff.value(T, y), outvalue, results) +end replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) # a perturbation riding on an argument ends up in the derivative, so `D` has to be able to hold it @@ -288,16 +290,11 @@ end end # an enclosing differentiation's tag is constant in our arguments, while one nested inside -# `f` buries our partial; a widened element type is a `Union`, so every member is checked +# `f` buries our partial @inline function checktags(::Type{T}, ::Type{E}) where {T, E} - if E isa Union - checktags(T, E.a) - checktags(T, E.b) - else - S = ForwardDiff.tagtype(E) - if S !== Nothing && !ForwardDiff.:≺(S, T) - throw(ForwardDiff.DualMismatchError(T, S)) - end + S = ForwardDiff.tagtype(E) + if S !== Nothing && !ForwardDiff.:≺(S, T) + throw(ForwardDiff.DualMismatchError(T, S)) end return nothing end @@ -322,7 +319,9 @@ end partials = reversepartials(results, input) foreach(targets) do (p, k, bound) x = getat(input, p) - istracked(x) && _br_add_to_deriv!(typeof(tag), x, k, output_deriv, partials, bound) + if istracked(x) + _br_add_to_deriv!(typeof(tag), x, k, output_deriv, partials, bound) + end end unseed!(output) return nothing diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index e20ee4aa..863df5fe 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -890,4 +890,12 @@ for i in eachindex(trs) @test samefields(trs[i], track(varr[i], tp2)) end +################### +# Pretty Printing # +################### + +@test occursin(r"^TrackedReal<\w{3}>\(1\.0, 0\.0, ---, ---\)$", repr(TrackedReal(1.0, 0.0))) +ta = TrackedArray([1.0, 2.0], [0.0, 3.0], InstructionTape()) +@test occursin(r"^TrackedReal<\w{3}>\(2\.0, 3\.0, \w{3}, 2, \w{3}\)$", repr(ta[2])) + end diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 19bb28c7..583987e6 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -285,8 +285,8 @@ end x -> sum(ifelse.(x .> 0, x, 0.0)) + sum(x), ) tape = ReverseDiff.GradientTape(f, [-1.0, -2.0]) - @test ReverseDiff.gradient!(tape, [1.0, 2.0]) == [2.0, 2.0] - @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [1.0, 2.0]) == [2.0, 2.0] + @test ReverseDiff.gradient!(tape, [1.5, 2.5]) == [2.0, 2.0] + @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [1.5, 2.5]) == [2.0, 2.0] end @test relu0.(track([-1.0, -2.0], InstructionTape())) isa TrackedArray @@ -327,7 +327,7 @@ end @test_throws ForwardDiff.DualMismatchError ReverseDiff.gradient( x -> sum((t -> t > 1.5 ? inner(t) : t).(x)), [1.0, 2.0] ) - # and next to a constant branch, which is not + # and next to a constant branch, which widens the results to `Real` @test_throws ForwardDiff.DualMismatchError ReverseDiff.gradient( x -> sum((t -> t > 1.5 ? inner(t) : 1.0).(x)), [1.0, 2.0] ) @@ -515,8 +515,13 @@ end @test length(take_recorded!(tp)) == 1 # an abstract element type is checked element by element, skipping non-numbers - g = x -> sum(x .* first.(((1.0, nothing), (2, nothing), (3.0, nothing)))) + b = ((1.0, nothing), (2.0, missing), (3.0, nothing)) + g = x -> sum(x .* (p -> first(p)::Float64).(b)) @test ReverseDiff.gradient(g, a) == [1.0, 2.0, 3.0] + + tp = InstructionTape() + track(copy(a), tp) .* (p -> first(p)::Float64).(b) + @test length(take_recorded!(tp)) == 1 end @testset "a `TrackedReal` without a tape is a constant" begin From adf6bb06472e582d395e3441f4efa5db4b924c78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 00:50:44 +0200 Subject: [PATCH 29/36] Bound a widened broadcast's replayed values by `Real` `ForwardDiff.valtype` unwraps only a concrete `Dual`, so a widened inferred element type (`Any` on Julia 1.10, or a `Dual` with free parameters) became the value buffer's element type, which rejected the values and hid the `DualMismatchError` of a nested tag. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 0e379320..ef8b828d 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -239,8 +239,10 @@ end # the cache carries what the replay needs: `df` to recompute the stored partials, or, where # they are known already, `vf` for the values alone function replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} - outvalue = similar(results, ForwardDiff.valtype(T, eltype(results))) - return df, map!(y -> ForwardDiff.value(T, y), outvalue, results) + E = eltype(results) + # a widened `E`, such as `Any` or a `Dual` with free parameters, bounds the values by `Real` + V = isconcretetype(E) ? ForwardDiff.valtype(T, E) : Real + return df, map!(y -> ForwardDiff.value(T, y), similar(results, V), results) end replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) From 6aca0607afe295cd3777c39e15769e79c40e28c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 09:17:18 +0200 Subject: [PATCH 30/36] Promote a broadcast's inferred result union instead of widening it `Broadcast.combine_eltypes` widens `Union{Dual, Float64}` to `Real`, so a function with a constant branch, such as `ifelse.(x .> 0, x, 0.0)`, cached boxed results and tracked `Real` values. Promoting the union keeps concrete `Dual`s and `Float64` values. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 6 +++++- test/derivatives/BroadcastTests.jl | 6 ++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index ef8b828d..5dd22a35 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -301,10 +301,14 @@ end return nothing end +# promoted instead of widened like `Broadcast.combine_eltypes`, so constant branches become `Dual`s +promoteunion(::Type{E}) where {E} = E +promoteunion(E::Union) = promote_type(promoteunion(E.a), promoteunion(E.b)) + # inferred, not read off the results, since a replay can take other branches and writes into # the results @inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, vals) where {T} - E = Broadcast.combine_eltypes(df, vals) + E = promoteunion(Base.promote_op(df, map(eltype, vals)...)) if typeintersect(E, Dual{T}) === Union{} checktags(T, E) return results diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 583987e6..040608c7 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -274,7 +274,7 @@ end tp = InstructionTape() y = relu.(track(copy(a), tp)) - @test y isa TrackedArray + @test y isa TrackedArray{Float64} @test value(y) == [0.0, 2.0] @test ReverseDiff.gradient(x -> sum(relu.(x)), a) == [0.0, 1.0] @@ -288,7 +288,9 @@ end @test ReverseDiff.gradient!(tape, [1.5, 2.5]) == [2.0, 2.0] @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [1.5, 2.5]) == [2.0, 2.0] end - @test relu0.(track([-1.0, -2.0], InstructionTape())) isa TrackedArray + x = track([-1.0, -2.0], InstructionTape()) + @test relu0.(x) isa TrackedArray{Float64} + @test ifelse.(x .> 0, x, 0.0) isa TrackedArray{Float64} # a static array whose widened element type is not isbits tape = ReverseDiff.GradientTape(x -> sum(relu.(x)), MVector(-1.0, 2.0)) From 0f323659e4d222c93affed17640f24af20e07cba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 09:28:48 +0200 Subject: [PATCH 31/36] Say what handles mismatched shapes in the propagation header Co-Authored-By: Claude Opus 5.5 --- src/derivatives/propagation.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index 68205788..b447376e 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -5,7 +5,7 @@ amount scaled by the output's derivative(s). Sometimes, extra partial informatio form of normal scalars/arrays, or `Dual` numbers) needs to accounted for as well. Often, a function's input and output values are not similarly shaped. To account for these -cases, `diffresult` and `contract` versions of the propagation functions have been implemented. +cases, the propagation functions take an index bound that clamps the output's indices to the input's. Accumulations into a scalar derivative go through `sum`, whose pairwise reduction is both more accurate and faster than a sequential loop. From 66b3f1cd7fc4a37de68db56a96be9b9d7ff81fd8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 10:08:15 +0200 Subject: [PATCH 32/36] Sharpen the broadcast tests and say why a broadcast target may be untracked Use asymmetric values for the tracked scalar factor, let the enclosing and the held perturbation reach the gradient, pin the nested-broadcast error message, read `value` without tracing, and test tracked values hidden in a closure or a `Ref`. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 1 + test/derivatives/BroadcastTests.jl | 29 ++++++++++++++++++++++------- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 5dd22a35..d77a3be5 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -325,6 +325,7 @@ end partials = reversepartials(results, input) foreach(targets) do (p, k, bound) x = getat(input, p) + # `capture` stores a `TrackedReal` without a tape as its value if istracked(x) _br_add_to_deriv!(typeof(tag), x, k, output_deriv, partials, bound) end diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 040608c7..c96374cb 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -212,7 +212,7 @@ end # an untracked container of tracked elements is a valid destination dest = Vector{TrackedReal{Float64, Float64, Nothing}}(undef, 4) dest .= exp.(x) - @test value.(dest) ≈ exp.(a) + @test value(dest) ≈ exp.(a) # the elements carry the tape, so it must not be dropped on the way in @test all(d -> tape(d) === tp, dest) end @@ -338,10 +338,10 @@ end xs = [1.0, 3.0] m(a) = sum( ReverseDiff.gradient( - x -> sum(ifelse.(x .> 2, x, a)), xs, ReverseDiff.GradientConfig(xs, typeof(a)) + x -> sum(ifelse.(x .> 2, x, a) .* a), xs, ReverseDiff.GradientConfig(xs, typeof(a)) ) ) - @test ForwardDiff.derivative(m, 1.0) == 0.0 + @test ForwardDiff.derivative(m, 1.0) == 1.0 end @testset "a perturbation the derivative cannot hold is rejected (#67, #168)" begin @@ -363,8 +363,12 @@ end @test_throws ArgumentError(msg3) ForwardDiff.derivative(g3, 3.0) # a tape whose derivatives are themselves `Dual`s can hold it, so it is left alone - h(a) = ReverseDiff.gradient(x -> sum(x .* a), [ForwardDiff.Dual(1.0, 0.0)]) - @test h(3.0) == [ForwardDiff.Dual(3.0, 0.0)] + h(a) = sum( + ReverseDiff.gradient( + x -> sum(x .* a), [1.0, 2.0], ReverseDiff.GradientConfig([1.0, 2.0], typeof(a)) + ) + ) + @test ForwardDiff.derivative(h, 3.0) == 2.0 # also when the tag without the perturbation was created first ReverseDiff.gradient(x -> sum(atan.(x, 2.0)), [1.0, 2.0]) @@ -409,7 +413,9 @@ end # also when the perturbation comes from an enclosing broadcast f(x) = sum(broadcast(a -> ReverseDiff.gradient(y -> sum(y .* a), [1.0])[1], x)) - @test_throws ArgumentError ReverseDiff.gradient(f, [2.0, 3.0]) + @test_throws "carries a perturbation that a derivative of type Float64 cannot hold" ReverseDiff.gradient( + f, [2.0, 3.0] + ) end @testset "zero-dimensional arrays (#265)" begin @@ -535,6 +541,14 @@ end end end +@testset "a tracked value hidden from `∇broadcast` takes the scalar rules" begin + u = [2.0, 3.0, 5.0] + expected = [2 * u[1] + sum(u[2:end]), u[1], u[1]] + + @test ReverseDiff.gradient(t -> sum((s -> s * t[1]).(t)), u) == expected + @test ReverseDiff.gradient(t -> sum(Ref(t[1]) .* t), u) == expected +end + @testset "`NotTracked`" begin f = ReverseDiff.NotTracked(t -> 2t) a = [1.0, 2.0, 3.0] @@ -614,7 +628,8 @@ end # a tracked scalar collects the whole output instead of one element of it @test ReverseDiff.gradient(t -> sum(t[1] .+ w), v) == [3.0, 0.0, 0.0] - @test ReverseDiff.gradient(t -> sum(t .* t[1]), v) == [2 * v[1] + sum(v[2:end]), 1.0, 1.0] + u = [2.0, 3.0, 5.0] + @test ReverseDiff.gradient(t -> sum(t .* t[1]), u) == [2 * u[1] + sum(u[2:end]), u[1], u[1]] # an argument `splitargs` holds back is not one a partial may name @test ReverseDiff.gradient(t -> sum(Ref(2.0) .* t), v) == fill(2.0, 3) From 341c92138862372b16ee2b94b8cef4eef430e124 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 15:02:21 +0200 Subject: [PATCH 33/36] Unwrap `NotTracked` before a scalar broadcast and tidy the broadcast tests A scalar broadcast passed a `NotTracked` argument to `f` still wrapped, which master unwrapped. Rename the eltype-changing `writable` to `replayable` next to its only caller, fix two comments, test through `@forward`/`@skip` and allocations instead of internals, skip the replay check for untracked results, and cover `.\` with a tracked denominator and the replay of a known partial naming an array of `TrackedReal`s. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 13 +++++++--- src/derivatives/propagation.jl | 2 +- src/tracked.jl | 4 --- test/derivatives/BroadcastTests.jl | 37 +++++++++++++++++----------- test/derivatives/ElementwiseTests.jl | 17 +++++++------ 5 files changed, 43 insertions(+), 30 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index d77a3be5..aea9b4d7 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -77,12 +77,13 @@ function remove_not_tracked(b::Broadcasted{style}) where {style} end function Base.copy(_bc::Broadcasted{<:TrackedStyle}) + bc = remove_not_tracked(_bc) # scalars take the scalar derivative rules. Ask the axes, since `TrackedStyle{Any}` carries # no dimension and `LinearAlgebra` may pass an uninstantiated `Broadcasted`. if axes(_bc) isa Tuple{} - return _bc[CartesianIndex()] + return bc[CartesianIndex()] end - flattened_bc = Base.Broadcast.flatten(remove_not_tracked(_bc)) + flattened_bc = Base.Broadcast.flatten(bc) f, args = flattened_bc.f, flattened_bc.args # only the arguments are seeded, not e.g. a closure's captures, and a `TrackedArray` # holds only `Real`s @@ -173,7 +174,7 @@ const RealOrArray = Union{Real, AbstractArray{<:Real}} # an argument read at the output's own index has to span the whole output ifsameshape(c::Contract, x, y) = (x isa Real || y isa Real || axes(x) == axes(y)) ? c : nothing -# partial with respect to argument `i`, or `nothing`. Requiring every argument to be +# partial with respect to the argument at the given position, or `nothing`. Requiring every argument to be # `RealOrArray` makes positions in `args` positions in the arguments `splitargs` keeps. knownpartial(f, ::Val, args) = nothing @@ -305,6 +306,10 @@ end promoteunion(::Type{E}) where {E} = E promoteunion(E::Union) = promote_type(promoteunion(E.a), promoteunion(E.b)) +# the replay writes values of type `E` into the results, which a static array may reject +replayable(results::AbstractArray, ::Type{E}) where {E} = convert(AbstractArray{E}, results) +replayable(results::StaticArray, ::Type{E}) where {E} = copyto!(similar(results, E), results) + # inferred, not read off the results, since a replay can take other branches and writes into # the results @inline function trackresults(::Type{T}, results::AbstractArray, df, vf, targs, vals) where {T} @@ -313,7 +318,7 @@ promoteunion(E::Union) = promote_type(promoteunion(E.a), promoteunion(E.b)) checktags(T, E) return results else - return recordresults(T, writable(results, E), df, vf, targs, vals) + return recordresults(T, replayable(results, E), df, vf, targs, vals) end end diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index b447376e..591fa4bd 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -5,7 +5,7 @@ amount scaled by the output's derivative(s). Sometimes, extra partial informatio form of normal scalars/arrays, or `Dual` numbers) needs to accounted for as well. Often, a function's input and output values are not similarly shaped. To account for these -cases, the propagation functions take an index bound that clamps the output's indices to the input's. +cases, the propagation functions take an index bound that clamps indices into the smaller array. Accumulations into a scalar derivative go through `sum`, whose pairwise reduction is both more accurate and faster than a sequential loop. diff --git a/src/tracked.jl b/src/tracked.jl index 4487a4aa..23a58363 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -552,10 +552,6 @@ writable(x::SArray) = copyto!(similar(x), x) writable(x::Adjoint) = adjoint(writable(parent(x))) writable(x::Transpose) = transpose(writable(parent(x))) -# the replay writes into broadcast results, which a static array may reject -writable(results::AbstractArray, ::Type{E}) where {E} = convert(AbstractArray{E}, results) -writable(results::StaticArray, ::Type{E}) where {E} = copyto!(similar(results, E), results) - # every forward pass writes into the value buffer, which a `TrackedArray` rejects # TODO: self-nesting, and the `value`/`track` unwrapping elsewhere, risk perturbation confusion (#45) track(x::TrackedArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = track(collect(x), D, tp) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index c96374cb..63c70b90 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -36,6 +36,14 @@ struct Scale s::Float64 end +# bytes allocated by recording `f` +function recordbytes(f, args...) + tp = InstructionTape() + targs = map(a -> track(a, tp), args) + f(targs...) + return @allocated f(targs...) +end + ############################################################################################ @testset "`BroadcastStyle` tracks dimensionality" begin @@ -140,14 +148,14 @@ end @testset "`@forward` and `@skip` wrappers do not force the fallback path" begin tp = InstructionTape() x = track(rand(3, 3), tp) - y = broadcast(ReverseDiff.ForwardOptimize(exp), x) + y = ReverseDiff.@forward(exp).(x) @test y isa TrackedArray @test length(take_recorded!(tp)) == 1 # `@skip` results are untracked, as with `map` and scalars tp = InstructionTape() x = track(rand(3, 3), tp) - y = broadcast(ReverseDiff.SkipOptimize(exp), x) + y = ReverseDiff.@skip(exp).(x) @test y isa Matrix{Float64} @test y == exp.(value(x)) @test isempty(take_recorded!(tp)) @@ -159,7 +167,7 @@ end @test ReverseDiff.hessian(f, a) ≈ zeros(3, 3) # arrays too - s = ReverseDiff.SkipOptimize(v -> sum(exp, v)) + s = ReverseDiff.@skip(v -> sum(exp, v)) @test ReverseDiff.gradient(v -> sum(v) * s(v), a) ≈ fill(sum(exp, a), 3) @test ReverseDiff.hessian(v -> sum(v) * s(v), a) == zeros(3, 3) end @@ -298,10 +306,10 @@ end end @testset "the cached partials have a concrete element type" begin - # the inferred element type decides what is recorded - tp = InstructionTape() - track(rand(3), tp) .^ 2 - @test isconcretetype(eltype(first(only(take_recorded!(tp)).cache))) + # one 16 B `Dual` per element besides the output's value and derivative, where boxed + # partials would take far more + n = 10_000 + @test recordbytes(t -> t .^ 2, rand(n)) < 40n end @testset "a non-`Real` result keeps its partials" begin @@ -562,19 +570,14 @@ end @test ReverseDiff.gradient(x -> sum(g.(x, ReverseDiff.NotTracked(Scale(2.0)))), a) == fill(2.0, 3) @test ReverseDiff.gradient(x -> sum(g.(x, ReverseDiff.NotTracked(Scale.(a)))), a) == a + @test ReverseDiff.gradient(x -> g.(x[1], ReverseDiff.NotTracked(Scale(2.0))), a) == + [2.0, 0.0, 0.0] end @testset "a partial known in closed form stores nothing" begin n = 10_000 v, w = rand(n), rand(n) - function recordbytes(f, args...) - tp = InstructionTape() - targs = map(a -> track(a, tp), args) - f(targs...) - return @allocated f(targs...) - end - # bytes allocated by the whole recording: the output's value and derivative take 8 B per # element each, and one `Dual` per element would add at least 16 B more for f in ( @@ -620,6 +623,7 @@ end # `-x/y^2` is no argument of the broadcast, so it is read off a `Dual` @test ReverseDiff.gradient(t -> sum(w ./ t), v) ≈ -w ./ v .^ 2 + @test ReverseDiff.gradient(t -> sum(t .\ w), v) ≈ -w ./ v .^ 2 # an outer broadcast still clamps the reverse pass to the argument's own shape @test ReverseDiff.gradient(t -> sum(t .+ w'), v) == fill(3.0, 3) @@ -658,6 +662,11 @@ end tape2 = ReverseDiff.GradientTape(t -> sum(t .* t), [1.0, 2.0, 3.0]) @test ReverseDiff.gradient!(tape2, x) == 2 .* x @test ReverseDiff.gradient!(ReverseDiff.compile(tape2), x) == 2 .* x + + # also when the argument is an array of `TrackedReal`s + tape3 = ReverseDiff.GradientTape(t -> sum(t .* [t[i] for i in 1:3]), [1.0, 2.0, 3.0]) + @test ReverseDiff.gradient!(tape3, x) == 2 .* x + @test ReverseDiff.gradient!(ReverseDiff.compile(tape3), x) == 2 .* x end end # module diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index da082efe..b03d38a3 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -54,7 +54,8 @@ function test_elementwise(f, fopt, x, tp) yt = broadcast(fopt, xt) @test yt == y # a function that returns no `Dual` has no derivative and is left off the tape - tracked = f(ForwardDiff.Dual(first(x), 1.0)) isa ForwardDiff.Dual + T = typeof(ForwardDiff.Tag(f, eltype(x))) + tracked = f(ForwardDiff.Dual{T}(first(x), one(eltype(x)))) isa ForwardDiff.Dual @test (yt isa ReverseDiff.TrackedArray) == tracked ReverseDiff.finish!(tp) @test length(tp) == tracked @@ -64,12 +65,14 @@ function test_elementwise(f, fopt, x, tp) ReverseDiff.seeded_reverse_pass!(out, yt, xt, tp) test_approx(out, ForwardDiff.jacobian(z -> broadcast(f, z), x); nans = true) - # forward - x2 = x .- offset - ReverseDiff.value!(xt, x2) - ReverseDiff.forward_pass!(tp) - @test yt == broadcast(f, x2) - ReverseDiff.value!(xt, x) + # forward, which cannot reach an untracked result (#313) + if tracked + x2 = x .- offset + ReverseDiff.value!(xt, x2) + ReverseDiff.forward_pass!(tp) + @test yt == broadcast(f, x2) + ReverseDiff.value!(xt, x) + end empty!(tp) return nothing From f9ca5ec257785d84e4b5b2dfd8c8ff3f7eca000e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 16:21:41 +0200 Subject: [PATCH 34/36] Let `setindex!` reject `TrackedArray` broadcast destinations and fill test gaps A `.=` into a `TrackedArray` now fails in `setindex!`, whose error carries the hint, rather than in a `materialize!` override that wrappers bypassed anyway. The elementwise tests decide whether a result is tracked by inference, as the implementation does, and known partials are tested with a variadic `+` and at second order. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 7 +------ src/derivatives/elementwise.jl | 1 + src/tracked.jl | 2 +- test/derivatives/BroadcastTests.jl | 13 ++++++++----- test/derivatives/ElementwiseTests.jl | 5 +++-- 5 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index aea9b4d7..3274e3fe 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -27,8 +27,8 @@ mayhidetracked(b::ForwardOptimize) = mayhidetracked(b.f) mayhidetracked(b::SkipOptimize) = mayhidetracked(b.f) mayhidetracked(b::Broadcasted) = mayhidetracked(b.f) || any(mayhidetracked, b.args) -# nothing inside an argument is seeded, so ask about the element type, not the container _mayhidetracked(::Type{<:NotTracked}) = false +# nothing inside an argument is seeded, so ask about the element type, not the container _mayhidetracked(::Type{<:AbstractArray{F}}) where {F} = _mayhidetracked(F) # a type that is not concrete, such as `Type{T}` or an abstract type, may have fields _mayhidetracked(::Type{F}) where {F} = !isconcretetype(F) || fieldcount(F) > 0 @@ -98,11 +98,6 @@ function Base.copy(_bc::Broadcasted{<:TrackedStyle}) return copy(Broadcast.instantiate(Broadcasted{style}(f, elargs))) end -_no_tracked_dest() = throw(ArgumentError("`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead.")) - -# `.=` hands `copyto!` the source's style, not the destination's, so catch the destination here -Broadcast.materialize!(::TrackedArray, ::Broadcasted) = _no_tracked_dest() - getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D getouttype(::Any) = Union{} diff --git a/src/derivatives/elementwise.jl b/src/derivatives/elementwise.jl index 331e67b2..e7346087 100644 --- a/src/derivatives/elementwise.jl +++ b/src/derivatives/elementwise.jl @@ -172,6 +172,7 @@ end output = instruction.output output_deriv = deriv(output) results = first(instruction.cache) + # a `DiffResult` holds partials ForwardDiff already extracted, so there is no tag to check if istracked(input) diffresult_increment_deriv!(Nothing, input, output_deriv, results, 1) else diff --git a/src/tracked.jl b/src/tracked.jl index 23a58363..3c9af0ad 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -424,7 +424,7 @@ Base.@propagate_inbounds function Base.view(t::TrackedArray, inds::ViewIndex...) return TrackedArray(v, view(deriv(t), inds...), tape(t)) end -Base.setindex!(t::TrackedArray, args...) = error("TrackedArrays do not support setindex!") +Base.setindex!(t::TrackedArray, args...) = throw(ArgumentError("`TrackedArray`s do not support `setindex!`, so they cannot be assigned to or used as a broadcast destination. Use `y = f.(x)` instead.")) Base.IndexStyle(::TrackedArray) = IndexLinear() diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 63c70b90..0caac1c7 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -205,17 +205,15 @@ end @testset "`TrackedArray`s are rejected as broadcast destinations" begin a = rand(4) - msg = "`TrackedArray`s do not support `setindex!` and cannot be used as a broadcast destination. Use `y = f.(x)` instead." + msg = "`TrackedArray`s do not support `setindex!`, so they cannot be assigned to or used as a broadcast destination. Use `y = f.(x)` instead." tp = InstructionTape() x = track(copy(a), tp) + @test_throws ArgumentError(msg) x[1] = 1.0 @test_throws ArgumentError(msg) track(zeros(4), tp) .= exp.(x) @test_throws ArgumentError(msg) track(zeros(4), tp) .= a @test_throws ArgumentError(msg) track(fill(1.0), tp) .= 2.0 - # a foreign style brings its own `copyto!` @test_throws ArgumentError(msg) track(zeros(2), tp) .= SVector(1.0, 2.0) - # nothing may be recorded before the failure - @test isempty(take_recorded!(tp)) # an untracked container of tracked elements is a valid destination dest = Vector{TrackedReal{Float64, Float64, Nothing}}(undef, 4) @@ -603,7 +601,7 @@ end @test recordbytes(t -> w ./ t, v) > 32n end -@testset "a known partial differentiates as the general path does" begin +@testset "known partials give the right gradients" begin v, w = [1.0, 2.0, 3.0], [4.0, 5.0, 6.0] @test ReverseDiff.gradient(t -> sum(t .* 2.0), v) == fill(2.0, 3) @@ -628,6 +626,7 @@ end # an outer broadcast still clamps the reverse pass to the argument's own shape @test ReverseDiff.gradient(t -> sum(t .+ w'), v) == fill(3.0, 3) @test ReverseDiff.gradient(t -> sum(t .* w'), v) == fill(sum(w), 3) + @test ReverseDiff.gradient(t -> sum(broadcast(+, t, w', 1.0)), v) == fill(3.0, 3) @test ReverseDiff.gradient(t -> sum(t .* 2.0), Float64[]) == Float64[] # a tracked scalar collects the whole output instead of one element of it @@ -642,6 +641,10 @@ end ga, gb = ReverseDiff.gradient((a, b) -> sum(a .- b), (v, w)) @test ga == fill(1.0, 3) @test gb == fill(-1.0, 3) + + # the argument a partial names is itself tracked at second order + @test ReverseDiff.hessian(t -> sum(t .* t), v) == diagm(fill(2.0, 3)) + @test ReverseDiff.hessian(t -> sum(t .* w), v) == zeros(3, 3) end @testset "a known partial does not overflow on a subnormal divisor" begin diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index b03d38a3..badcc037 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -53,9 +53,10 @@ function test_elementwise(f, fopt, x, tp) # record yt = broadcast(fopt, xt) @test yt == y - # a function that returns no `Dual` has no derivative and is left off the tape + # a function inferred to return no `Dual` has no derivative and is left off the tape T = typeof(ForwardDiff.Tag(f, eltype(x))) - tracked = f(ForwardDiff.Dual{T}(first(x), one(eltype(x)))) isa ForwardDiff.Dual + R = Base.promote_op(f, ForwardDiff.Dual{T, eltype(x), 1}) + tracked = typeintersect(R, ForwardDiff.Dual) !== Union{} @test (yt isa ReverseDiff.TrackedArray) == tracked ReverseDiff.finish!(tp) @test length(tp) == tracked From dd47fc9e8b07c79c43cf10ea307ed39f516e2300 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 5 Oct 2026 23:03:10 +0200 Subject: [PATCH 35/36] Take the scalar rules for perturbations the derivative cannot hold A broadcast argument whose perturbation the derivative type cannot hold, such as a `Dual` created inside the differentiated function, now takes the scalar rules instead of throwing. Reverse-over-forward through a broadcast therefore works, and forward-over-reverse raises the #45 error. Also raise a descriptive error for output elements that are not numbers, cover the abstract-inference path of the reverse pass, and name the constant DiffTests function explicitly. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/broadcast.jl | 54 +++++++--------------- src/derivatives/propagation.jl | 9 ++-- src/tracked.jl | 3 ++ test/api/JacobianTests.jl | 3 ++ test/derivatives/BroadcastTests.jl | 68 ++++++++++++++++------------ test/derivatives/ElementwiseTests.jl | 8 ++-- 6 files changed, 71 insertions(+), 74 deletions(-) diff --git a/src/derivatives/broadcast.jl b/src/derivatives/broadcast.jl index 3274e3fe..cacad7ec 100644 --- a/src/derivatives/broadcast.jl +++ b/src/derivatives/broadcast.jl @@ -28,7 +28,7 @@ mayhidetracked(b::SkipOptimize) = mayhidetracked(b.f) mayhidetracked(b::Broadcasted) = mayhidetracked(b.f) || any(mayhidetracked, b.args) _mayhidetracked(::Type{<:NotTracked}) = false -# nothing inside an argument is seeded, so ask about the element type, not the container +# elements that are not `Real` are not seeded, so ask about the element type, not the container _mayhidetracked(::Type{<:AbstractArray{F}}) where {F} = _mayhidetracked(F) # a type that is not concrete, such as `Type{T}` or an abstract type, may have fields _mayhidetracked(::Type{F}) where {F} = !isconcretetype(F) || fieldcount(F) > 0 @@ -85,11 +85,12 @@ function Base.copy(_bc::Broadcasted{<:TrackedStyle}) end flattened_bc = Base.Broadcast.flatten(bc) f, args = flattened_bc.f, flattened_bc.args - # only the arguments are seeded, not e.g. a closure's captures, and a `TrackedArray` - # holds only `Real`s + # only the arguments are seeded, not e.g. a closure's captures, a `TrackedArray` holds + # only `Real`s, and the derivatives have to hold the arguments' perturbations if !mayhidetracked(_bc) vals = map(value, args) - if Broadcast.combine_eltypes(f, vals) <: Real + D = mapreduce(getouttype, promote_type, args) + if Broadcast.combine_eltypes(f, vals) <: Real && all(v -> holdsperturbations(D, v), vals) return ∇broadcast(f, args, vals) end end @@ -102,6 +103,18 @@ getouttype(::TrackedReal{<:Any, D}) where {D} = D getouttype(::AbstractArray{<:TrackedReal{<:Any, D}}) where {D} = D getouttype(::Any) = Union{} +# a perturbation riding on an argument ends up in the derivative, so `D` has to hold it. An +# abstract element type such as `Real` can hide one. +holdsperturbations(::Type, _) = true +function holdsperturbations(::Type{D}, v::Union{Real, AbstractArray, Tuple}) where {D} + E = eltype(v) + if isconcretetype(E) + return ForwardDiff.tagtype(E) === Nothing || promote_type(D, E) <: D + else + return all(x -> holdsperturbations(D, x), v) + end +end + deref(x) = x deref(x::Base.RefValue) = x[] @@ -242,41 +255,8 @@ function replaycache(::Type{T}, results::AbstractArray, df, _, _) where {T} end replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) -# a perturbation riding on an argument ends up in the derivative, so `D` has to be able to hold it -@inline function checkargtags(::Type{D}, ::Type{E}) where {D, E} - if ForwardDiff.tagtype(E) !== Nothing && !(promote_type(D, E) <: D) - throw( - ArgumentError( - LazyString( - "a broadcast argument with element type ", E, - " carries a perturbation that a derivative of type ", D, - " cannot hold" - ) - ) - ) - end - return nothing -end - -# an abstract element type such as `Real` can hide a perturbation -function checkargvalues(::Type{D}, v) where {D} - if isconcretetype(eltype(v)) - checkargtags(D, eltype(v)) - else - for x in v - if x isa Dual - checkargtags(D, typeof(x)) - elseif x isa Union{Tuple, AbstractArray} - checkargvalues(D, x) - end - end - end - return nothing -end - @inline function recordresults(::Type{T}, results, df, vf, targs, vals) where {T} D = mapreduce(getouttype, promote_type, targs) - foreach(v -> checkargvalues(D, v), vals) replayf, outvalue = replaycache(T, results, df, vf, vals) tp = tape(targs...) out = track(outvalue, D, tp) diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index 591fa4bd..43a9be5b 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -176,10 +176,11 @@ function contract_increment_deriv!( input::TrackedReal, x::AbstractArray, e::Contract, args::Tuple, ::Nothing ) - isempty(x) && return nothing - pull_deriv!(input) - input.deriv += sum(i -> _contract(e, x[i], i, args), eachindex(x)) - push_deriv!(input) + if !isempty(x) + pull_deriv!(input) + input.deriv += sum(i -> _contract(e, x[i], i, args), eachindex(x)) + push_deriv!(input) + end return nothing end diff --git a/src/tracked.jl b/src/tracked.jl index 3c9af0ad..91b0dccd 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -216,6 +216,9 @@ function checkoutput(::Dual{<:Any, <:TrackedReal}) throw(ArgumentError("ForwardDiff cannot differentiate through ReverseDiff (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)")) end +function seed!(x) + throw(ArgumentError(LazyString("an output must hold real numbers, but holds an element of type ", typeof(x)))) +end seed!(x::Real) = checkoutput(x) seed!(t::TrackedReal) = (t.deriv = one(derivtype(t)); push_deriv!(t); nothing) seed!(t::TrackedArray, i) = (t.deriv[i] = one(derivtype(t)); nothing) diff --git a/test/api/JacobianTests.jl b/test/api/JacobianTests.jl index 01af815a..8c09f88f 100644 --- a/test/api/JacobianTests.jl +++ b/test/api/JacobianTests.jl @@ -344,6 +344,9 @@ end @testset "outputs whose element type is not `Real`" begin @test ReverseDiff.jacobian(x -> Any[2 * x[1], 3 * x[2]], [1.0, 2.0]) == [2.0 0.0; 0.0 3.0] @test ReverseDiff.jacobian(x -> [1.0, 2.0], rand(3)) == zeros(2, 3) + + # an element that is not a number has no single derivative to seed + @test_throws ArgumentError ReverseDiff.jacobian(x -> [x, x], [1.0, 2.0]) end end # module diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 0caac1c7..295d4bff 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -31,6 +31,10 @@ Base.similar(::Broadcasted{ForeignStyle}, ::Type{T}, axs) where {T} = similar(Ar # at top level, since a local function that captures itself takes the scalar rules inner(t) = ForwardDiff.Dual{typeof(ForwardDiff.Tag(inner, typeof(t)))}(t, one(t)) +# a zero that inference only knows as `Real` +const REALZERO = Ref{Real}(0.0) +relureal(t) = t > 0 ? t : REALZERO[] + # a struct with fields, which `NotTracked` declares constant struct Scale s::Float64 @@ -270,6 +274,14 @@ end @test ReverseDiff.gradient(x -> sum((t -> 1.0).(x)), a) == zeros(3) @test ReverseDiff.gradient(x -> sum(oneunit.(x)), a) == zeros(3) + + # a function that returns no `Dual` is left untracked and off the tape + for f in (t -> 1.0, round, floor) + tp = InstructionTape() + y = f.(track(copy(a), tp)) + @test y isa Vector{Float64} + @test isempty(take_recorded!(tp)) + end end @testset "a type-unstable function keeps its partials" begin @@ -301,6 +313,11 @@ end # a static array whose widened element type is not isbits tape = ReverseDiff.GradientTape(x -> sum(relu.(x)), MVector(-1.0, 2.0)) @test ReverseDiff.gradient!(tape, MVector(1.0, 2.0)) == [1.0, 1.0] + + # inferred as `Real`, so the constant branch's results stay plain numbers + @test ReverseDiff.gradient(x -> sum(relureal.(x)), a) == [0.0, 1.0] + tape = ReverseDiff.GradientTape(x -> sum(relureal.(x)), [-1.0, -2.0]) + @test ReverseDiff.gradient!(tape, [1.5, 2.5]) == [1.0, 1.0] end @testset "the cached partials have a concrete element type" begin @@ -350,23 +367,23 @@ end @test ForwardDiff.derivative(m, 1.0) == 1.0 end -@testset "a perturbation the derivative cannot hold is rejected (#67, #168)" begin - # `a` reaches the broadcast as a `Dual`, but the tape's derivatives are `Float64` - g(a) = sum(ReverseDiff.gradient(x -> sum(x .* a), [1.0, 2.0])) - E = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g), Float64}, Float64, 1} - msg = "a broadcast argument with element type $E carries a perturbation that a derivative of type Float64 cannot hold" - - @test_throws ArgumentError(msg) ForwardDiff.derivative(g, 3.0) - - # an abstract element type is checked element by element - g2(a) = sum(ReverseDiff.gradient(x -> sum(x .* Union{Float64, typeof(a)}[1.0, a]), [1.0, 2.0])) - E2 = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g2), Float64}, Float64, 1} - msg2 = "a broadcast argument with element type $E2 carries a perturbation that a derivative of type Float64 cannot hold" - @test_throws ArgumentError(msg2) ForwardDiff.derivative(g2, 3.0) +@testset "a perturbation the derivative cannot hold takes the scalar rules (#67, #168)" begin + # the scalar rules bury the tracked value in the `Dual`, which must not give a zero derivative + msg = "ForwardDiff cannot differentiate through ReverseDiff (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)" + for g in ( + # `a` reaches the broadcast as a `Dual`, but the tape's derivatives are `Float64` + a -> sum(ReverseDiff.gradient(x -> sum(x .* a), [1.0, 2.0])), + # an abstract element type is checked element by element + a -> sum(ReverseDiff.gradient(x -> sum(x .* Union{Float64, typeof(a)}[1.0, a]), [1.0, 2.0])), + a -> sum(ReverseDiff.gradient(x -> sum(x .* Real[1.0, a]), [1.0, 2.0])), + a -> sum(ReverseDiff.gradient(x -> x[1] * a, [1.0, 2.0])), + a -> sum(ReverseDiff.jacobian(x -> [x[1] * a, x[2]], [1.0, 2.0])), + ) + @test_throws ArgumentError(msg) ForwardDiff.derivative(g, 3.0) + end + # a perturbation that never meets a tracked value leaves the derivative exact g3(a) = sum(ReverseDiff.gradient(x -> sum(ifelse.(x .> 5, x, Real[1.0, a])), [1.0, 2.0])) - E3 = ForwardDiff.Dual{ForwardDiff.Tag{typeof(g3), Float64}, Float64, 1} - msg3 = "a broadcast argument with element type $E3 carries a perturbation that a derivative of type Float64 cannot hold" - @test_throws ArgumentError(msg3) ForwardDiff.derivative(g3, 3.0) + @test ForwardDiff.derivative(g3, 3.0) == 0.0 # a tape whose derivatives are themselves `Dual`s can hold it, so it is left alone h(a) = sum( @@ -387,15 +404,6 @@ end xs = [1.0, 2.0] @test ForwardDiff.derivative(k, 3.0) ≈ sum(@. (xs^2 - 9) / (xs^2 + 9)^2) - # the scalar rules bury the tracked value in the `Dual`, which must not give a zero derivative - msg4 = "ForwardDiff cannot differentiate through ReverseDiff (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)" - for g4 in ( - a -> sum(ReverseDiff.gradient(x -> x[1] * a, [1.0, 2.0])), - a -> sum(ReverseDiff.gradient(x -> sum(x .* Real[1.0, a]), [1.0, 2.0])), - a -> sum(ReverseDiff.jacobian(x -> [x[1] * a, x[2]], [1.0, 2.0])), - ) - @test_throws ArgumentError(msg4) ForwardDiff.derivative(g4, 3.0) - end # storing the `Dual` as a tracked number would cut it off the tape msg5 = "this nesting of ForwardDiff and ReverseDiff is not supported: a `Dual` of tracked numbers cannot be converted to a tracked number (see https://github.com/JuliaDiff/ReverseDiff.jl/issues/45)" function g5(a) @@ -416,12 +424,16 @@ end x -> ForwardDiff.derivative(a -> sum([x[1] * a^2, x[2]]), 2.0), [1.0, 2.0] ) == [4.0, 0.0] + # and through a broadcast, also on a replay + r = x -> ForwardDiff.derivative(a -> sum(x .* a^2), 2.0) + @test ReverseDiff.gradient(r, [1.0, 2.0]) == [4.0, 4.0] + @test ReverseDiff.gradient!(ReverseDiff.GradientTape(r, [1.0, 2.0]), [3.0, 4.0]) == [4.0, 4.0] + @test ReverseDiff.gradient(x -> ForwardDiff.derivative(a -> sum(x * a^2), 2.0), [1.0, 2.0]) == + [4.0, 4.0] # also when the perturbation comes from an enclosing broadcast f(x) = sum(broadcast(a -> ReverseDiff.gradient(y -> sum(y .* a), [1.0])[1], x)) - @test_throws "carries a perturbation that a derivative of type Float64 cannot hold" ReverseDiff.gradient( - f, [2.0, 3.0] - ) + @test_throws ArgumentError(msg) ReverseDiff.gradient(f, [2.0, 3.0]) end @testset "zero-dimensional arrays (#265)" begin diff --git a/test/derivatives/ElementwiseTests.jl b/test/derivatives/ElementwiseTests.jl index badcc037..54486273 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -53,13 +53,11 @@ function test_elementwise(f, fopt, x, tp) # record yt = broadcast(fopt, xt) @test yt == y - # a function inferred to return no `Dual` has no derivative and is left off the tape - T = typeof(ForwardDiff.Tag(f, eltype(x))) - R = Base.promote_op(f, ForwardDiff.Dual{T, eltype(x), 1}) - tracked = typeintersect(R, ForwardDiff.Dual) !== Union{} + # a constant function has no derivative and is left off the tape + tracked = f !== DiffTests.num2num_4 @test (yt isa ReverseDiff.TrackedArray) == tracked ReverseDiff.finish!(tp) - @test length(tp) == tracked + @test length(tp) == (tracked ? 1 : 0) # reverse out = similar(y, (length(x), length(x))) From dd16d6589d5ea6f117d9e415c83da885b5d8d3c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 6 Oct 2026 13:48:23 +0200 Subject: [PATCH 36/36] Test the `besselj` `NaN` partial fused with a tracked factor `besselj.(nu, y)` alone already passed on master; fused with `.* y`, the untracked order's `NaN` partial shares a `Dual` with the tracked partials and gave a `NaN` gradient there. Co-Authored-By: Claude Opus 5.5 --- test/derivatives/BroadcastTests.jl | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/test/derivatives/BroadcastTests.jl b/test/derivatives/BroadcastTests.jl index 295d4bff..3f3e1e13 100644 --- a/test/derivatives/BroadcastTests.jl +++ b/test/derivatives/BroadcastTests.jl @@ -231,9 +231,12 @@ end # `DiffRules` gives `besselj` a `NaN` partial for its order, which is untracked here nu, z = [0.5, 1.5, 2.5], 0.7 # J_ν'(z) = (J_{ν-1}(z) - J_{ν+1}(z)) / 2 - expected = sum((besselj(n - 1, z) - besselj(n + 1, z)) / 2 for n in nu) + dbesselj(n, z) = (besselj(n - 1, z) - besselj(n + 1, z)) / 2 - @test ReverseDiff.gradient(y -> sum(besselj.(nu, y)), [z]) ≈ [expected] + @test ReverseDiff.gradient(y -> sum(besselj.(nu, y)), [z]) ≈ [sum(dbesselj(n, z) for n in nu)] + # fused with a tracked factor, the `NaN` shares a `Dual` with the tracked partials + @test ReverseDiff.gradient(y -> sum(besselj.(nu, y) .* y), [z]) ≈ + [sum(dbesselj(n, z) * z + besselj(n, z) for n in nu)] end @testset "an array of `TrackedReal`s is differentiated like a `TrackedArray`" begin