diff --git a/Project.toml b/Project.toml index 41bc0d21..2a35c6ce 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ReverseDiff" uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267" -version = "1.18.5" +version = "1.18.6" [deps] ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" diff --git a/src/api/utils.jl b/src/api/utils.jl index bb6fb33c..bd959ee9 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 60d9d38f..cacad7ec 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) @@ -17,46 +15,59 @@ 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(::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 - # `Type{T}` and abstract types; be conservative in that case. - hasfields = try - fieldcount(F) > 0 - catch - true - end - return :($hasfields) +# 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(::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) + +_mayhidetracked(::Type{<:NotTracked}) = false +# 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 + +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}() + +# `AbstractArrayStyle{Any}` carries `Any` as its dimension, which `max` cannot compare +_maxdim(M::Int, N::Int) = max(M, N) +_maxdim(_, _) = 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 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 -mayhavetracked(b) = istypeorclosure(b) -mayhavetracked(b::Type) = false -mayhavetracked(b::NotTracked) = false -mayhavetracked(b::Base.RefValue{<:NotTracked}) = false -mayhavetracked(b::Broadcasted) = mayhavetracked(b.f) || any(mayhavetracked, b.args) - -struct TrackedStyle <: BroadcastStyle end - -Broadcast.BroadcastStyle(::Type{<:Union{TrackedArray, TrackedReal}}) = TrackedStyle() -Broadcast.BroadcastStyle(::TrackedStyle, b::BroadcastStyle) = TrackedStyle() - -# 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)) - -function broadcast_rebuild(bc::Broadcasted) - return broadcasted(bc.f, broadcast_rebuild.(bc.args)...) -end - -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 +76,49 @@ 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) - elseif !mayhavetracked(bc) && outputisreal - # 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. - return Val(:reversediff) - else - # 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 - return Val(:fallback) - end -end -function Base.copy(_bc::Broadcasted{TrackedStyle}) +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()] + end 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...) - 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)) + # 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) + 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 + elargs = map(staticelements, args) + style = typeof(reduce(Broadcast.result_style, map(untrackedstyle, elargs))) + return copy(Broadcast.instantiate(Broadcasted{style}(f, elargs))) end 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{} +# 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[] -@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 @@ -138,8 +139,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)]...) @@ -148,157 +148,222 @@ 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) +@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. +@inline trackedslots(::Tuple{}) = (), (), Val(0) +@inline function trackedslots(args::Tuple) + slots, positions, n = trackedslots(Base.front(args)) + if istracked(last(args)) + k = incr(n) + return (slots..., k), (positions..., Val(length(args))), k else - _br_add_to_deriv!(input, output_deriv, results, bounds) + return (slots..., Val(0)), positions, n end - unseed!(output) - return nothing 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) -end +@inline getat(xs::Tuple, ::Val{i}) where {i} = xs[i] -@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) +# `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 -@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 +# `f`'s partials in closed form, one per tracked argument +struct KnownPartials{E <: Tuple} + entries::E end -## Tracker style broadcasting -## Good for broadcasting real numbers or arrays of non-tracked structs - -trim(x, Δ) = reshape(Δ, ntuple(i -> size(Δ, i), Val(ndims(x)))) - -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(Δ))))) +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 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 + +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(-, ()) + +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}) = + 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 + p = knownpartial(f, first(positions), args) + if p === nothing + return nothing + else + return (p, rest...) + end +end -unbroadcast(x::Number, Δ) = sum(Δ) -unbroadcast(x::Base.RefValue, _) = nothing +# marks the `Dual`s `∇broadcast` seeds, so `@skip` can drop them as it drops tracking +struct BroadcastTag{F} end -dual(x, p) = x -dual(x::Real, p) = Dual(x, p) +skipvalue(x::Dual{T}) where {T <: ForwardDiff.Tag{<:BroadcastTag}} = skipvalue(ForwardDiff.value(T, x)) -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 +# 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) + slots, positions, valP = trackedslots(targs) + # keyed on every argument's type, so an enclosing differentiation's `Dual` makes a newer tag + 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(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(typeof(tag), broadcast(df, vals...), df, vf, targs, vals) + else + return recordresults(typeof(tag), KnownPartials(entries), df, vf, targs, vals) + end 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]...) + +# 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} + 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 -@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) +replaycache(::Type, ::KnownPartials, _, vf, vals) = (vf, broadcast(vf, vals...)) + +@inline function recordresults(::Type{T}, results, df, vf, targs, vals) where {T} + D = mapreduce(getouttype, promote_type, targs) + 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, replayf, T(), map(tuple, positions, ntuple(Val, n), bounds)) + 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)...) +# an enclosing differentiation's tag is constant in our arguments, while one nested inside +# `f` buries our partial +@inline function checktags(::Type{T}, ::Type{E}) where {T, E} + S = ForwardDiff.tagtype(E) + if S !== Nothing && !ForwardDiff.:≺(S, T) + throw(ForwardDiff.DualMismatchError(T, S)) + end return nothing end -@noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(tracker_∇broadcast)}) +# 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)) + +# 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} + E = promoteunion(Base.promote_op(df, map(eltype, vals)...)) + if typeintersect(E, Dual{T}) === Union{} + checktags(T, E) + return results + else + return recordresults(T, replayable(results, E), df, vf, targs, vals) + 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, targets = instruction.cache + 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 + end unseed!(output) return nothing end -## Limited ReverseDiff broadcasting -## Efficient broadcasting for specific functions, e.g. +, - +struct PartialsWithArgs{E <: Tuple, A <: Tuple} + entries::E + args::A +end -@inline _materialize(f, args) = broadcast(f, args...) +# 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 +function reversepartials(p::KnownPartials, input) + args = map(e -> map(j -> value(getat(input, j)), e.args), p.entries) + return PartialsWithArgs(p.entries, args) +end -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, out_deriv, results, + bound::CartesianIndex + ) where {T} + if bound == CartesianIndex(size(out_deriv)) + return _increment_deriv!(T, x, out_deriv, results, slot) + else + 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) 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}, 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[k], bound...) + +@noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(∇broadcast)}) + input, output = instruction.input, instruction.output + results, replayf, tag, _ = instruction.cache + foreach(pull_value!, input) + _replay!(typeof(tag), value(output), results, replayf, 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 5caf1274..e7346087 100644 --- a/src/derivatives/elementwise.jl +++ b/src/derivatives/elementwise.jl @@ -1,32 +1,28 @@ -############################### -# SkipOptimized map/broadcast # -############################### +##################### +# SkipOptimized map # +##################### # dispatch # #----------# -for g in (:map, :broadcast), 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, :broadcast), 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) 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) 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,191 +30,49 @@ end # record # #--------# -for g in (:map, :broadcast) - @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, :broadcast) +@inline Base.map(f::SkipOptimize{F}, x::TrackedArray, 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 - 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 -end - -#################################### -# ForwardOptimized map!/broadcast! # -#################################### - -# 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 + @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)) 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) + @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 -################################## -# ForwardOptimized map/broadcast # -################################## +######################## +# ForwardOptimized map # +######################## # dispatch # #----------# -for g in (:map, :broadcast), (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 - # 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) - @inline Base.$(g)(f::typeof($M.$f), x::TrackedReal, y::TrackedArray) = $(g)(ForwardOptimize(f), x, y) - end + @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.$(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) + @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 @@ -227,119 +81,86 @@ end # record # #--------# -for g in (:map, :broadcast) - @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 + +# 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)) + end + results = map(df, value(x), value(y)) tp = tape(x) - out = track(DiffResults.value.(results), D, tp) - cache = (results, df, index_bound(x, out), nothing) - record!(tp, SpecialInstruction, $(g), x, out, cache) + 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 - 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(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} + if istracked(x) + return record_map(f, x, y, D) 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(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 - ForwardDiff.gradient!(result, s -> f.f(s[1], vy), SVector(vx)) - end - results = $(g)(df, value(x), value(y)) - tp = tape(x) - out = track(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::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(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 - - 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(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 + 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 = 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, map, (x, y), out, cache) + return out 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 +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)) + 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), (:broadcast!, :broadcast)) - @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 ################ @@ -351,324 +172,21 @@ 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!(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) + diffresult_increment_deriv!(Nothing, 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) - 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/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index c1f1ef56..1a898efa 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -129,10 +129,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)) @@ -163,14 +159,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 @@ -256,9 +252,9 @@ LinearAlgebra.:*(x::StaticMatrix, y::TrackedVector{V, D}) where {V, D} = record_ 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 ab9ca6f0..43a9be5b 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, `broadcast` and `reduce` versions of the propagation functions have been implemented. +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. @@ -95,142 +95,95 @@ 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) + +# 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 + +_elem(v::Real, i) = v +_elem(v::AbstractArray, i) = v[i] + +# `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!( - input::AbstractArray, x::AbstractArray, - results, p::Int - ) + ::Type{T}, input::AbstractArray, x::AbstractArray, + results::AbstractArray, p::Int + ) where {T} for i in eachindex(x, results) - increment_deriv!(input, x[i] * getpartial(results[i], p), i) + increment_deriv!(input, x[i] * getpartial(T, results[i], p), i) end return nothing end function diffresult_increment_deriv!( - input::AbstractArray, x::AbstractArray, - results, p::Int, bound::CartesianIndex - ) + ::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(r, p), min(bound, i)) + increment_deriv!(input, xi * getpartial(T, r, p), min(bound, i)) end return nothing end function diffresult_increment_deriv!( - input::TrackedReal, x::AbstractArray, - results, p::Int, ::Nothing - ) + ::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 -> x[i] * getpartial(results[i], p), inds) + input.deriv += sum(i -> x[i] * getpartial(T, results[i], p), inds) push_deriv!(input) return nothing end -############################## -# broadcast_increment_deriv! # -############################## +############################# +# contract_increment_deriv! # +############################# -# without partials # -#------------------# - -function broadcast_increment_deriv!( - input::AbstractArray, x::AbstractArray, - bound::CartesianIndex +function contract_increment_deriv!( + input::AbstractArray, x::AbstractArray, e::Contract, + args::Tuple ) - for (xi, i) in zip(x, CartesianIndices(size(x))) - increment_deriv!(input, xi, min(bound, 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)) + for i in eachindex(input, x) + increment_deriv!(input, _contract(e, x[i], i, args), 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) - isempty(inds) && return nothing - pull_deriv!(input) - input.deriv += sum(i -> op(x[i], partials[i]), 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)) +function contract_increment_deriv!( + input::AbstractArray, x::AbstractArray, e::Contract, + args::Tuple, bound::CartesianIndex + ) + for i in CartesianIndices(size(x)) + increment_deriv!(input, _contract(e, x[i], i, args), min(bound, i)) end return nothing end -############################## -# broadcast_decrement_deriv! # -############################## - -function broadcast_decrement_deriv!( - input::AbstractArray, x::AbstractArray, - bound::CartesianIndex +function contract_increment_deriv!( + input::TrackedReal, x::AbstractArray, e::Contract, + args::Tuple, ::Nothing ) - for (xi, i) in zip(x, CartesianIndices(size(x))) - decrement_deriv!(input, xi, min(bound, i)) + 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 -function broadcast_decrement_deriv!(input::TrackedReal, x::AbstractArray, ::Nothing) - isempty(x) && return nothing - pull_deriv!(input) - input.deriv -= sum(x) - push_deriv!(input) - return nothing -end - ############################## # reduction_increment_deriv! # ############################## diff --git a/src/macros.jl b/src/macros.jl index 083e83ee..5d8c1326 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -162,10 +162,11 @@ 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)) +skipvalue(x::AbstractArray{<: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 d9babb5f..7b592398 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -110,6 +110,9 @@ 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)) @inline deriv(t::TrackedArray) = t.deriv @inline deriv(t::TrackedReal) = t.deriv @@ -206,7 +209,17 @@ push_deriv!(x::AbstractArray) = (istracked(x) && foreach(push_deriv!, x); nothin # seed/unseed # #-------------# -seed!(x) = 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 + +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) seed!(x::AbstractArray, i) = seed!(x[i]) @@ -273,6 +286,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 @@ -292,6 +310,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} @@ -406,7 +427,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() @@ -533,6 +554,10 @@ writable(x::SArray) = copyto!(similar(x), x) writable(x::Adjoint) = adjoint(writable(parent(x))) writable(x::Transpose) = transpose(writable(parent(x))) +# 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) + track!(t::TrackedArray, x::AbstractArray) = (value!(t, x); unseed!(t); t) track!(t::TrackedReal, x::Real) = (value!(t, x); unseed!(t); t) @@ -551,9 +576,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) - 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/PrecisionTests.jl b/test/PrecisionTests.jl index 09c7cb90..6ba039a3 100644 --- a/test/PrecisionTests.jl +++ b/test/PrecisionTests.jl @@ -33,10 +33,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 19e90514..863df5fe 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -397,16 +397,11 @@ 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 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 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) @@ -424,16 +419,6 @@ ReverseDiff.deriv!(tr, d) @test ReverseDiff.value(tr) === v @test ReverseDiff.deriv(tr) === d -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 @@ -453,16 +438,6 @@ ReverseDiff.pull_deriv!(trs) @test ReverseDiff.deriv(trs[2]) === darr[1] @test ReverseDiff.deriv(trs[3]) === darr[2] -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 @@ -794,7 +769,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 # @@ -914,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/api/JacobianTests.jl b/test/api/JacobianTests.jl index eaff84ad..8c09f88f 100644 --- a/test/api/JacobianTests.jl +++ b/test/api/JacobianTests.jl @@ -341,4 +341,12 @@ 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) + + # 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 new file mode 100644 index 00000000..3f3e1e13 --- /dev/null +++ b/test/derivatives/BroadcastTests.jl @@ -0,0 +1,690 @@ +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) + +# 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 +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 + 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}()), + (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}()), + ) + @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}() + @test result_style(TrackedStyle{Any}(), BroadcastStyle(typeof(s))) === TrackedStyle{Any}() + + 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 + + # 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 StaticArray + @test z isa StaticArray + ReverseDiff.seed!(sum(y) + sum(z)) + 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 + +@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) + instrs1, instrs2 = take_recorded!(tp1), take_recorded!(tp2) + + @test value(x) == value(y) + @test length(instrs1) == length(instrs2) == 1 + @test instrs1[1].func === instrs2[1].func + end +end + +@testset "`@forward` and `@skip` wrappers do not force the fallback path" begin + tp = InstructionTape() + x = track(rand(3, 3), tp) + 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 = ReverseDiff.@skip(exp).(x) + @test y isa Matrix{Float64} + @test y == exp.(value(x)) + @test isempty(take_recorded!(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) + + # arrays too + 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 + +@testset "a type broadcast as a function does not force the fallback path" begin + tp = InstructionTape() + x = track(rand(3), tp) + + y = Real.(x) + instrs = take_recorded!(tp) + + @test y isa TrackedArray + @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 + +@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(take_recorded!(tp)) == 2 +end + +@testset "`TrackedArray`s are rejected as broadcast destinations" begin + a = rand(4) + 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 + @test_throws ArgumentError(msg) track(zeros(2), tp) .= SVector(1.0, 2.0) + + # 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 + dbesselj(n, z) = (besselj(n - 1, z) - besselj(n + 1, z)) / 2 + + @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 + 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 + s = sum(y) + ReverseDiff.finish!(tp) + + @test y isa TrackedArray + @test value(y) ≈ a .* b + # one instruction for the broadcast, one for `sum` + @test length(tp) == 2 + @test ReverseDiff.instructions(tp)[1].func === ReverseDiff.∇broadcast + + 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 +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) + + @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 + # 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{Float64} + @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.5, 2.5]) == [2.0, 2.0] + @test ReverseDiff.gradient!(ReverseDiff.compile(tape), [1.5, 2.5]) == [2.0, 2.0] + end + 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)) + @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 + # 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 + 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 + # 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] + ) + # 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] + ) + + # 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) .* a), xs, ReverseDiff.GradientConfig(xs, typeof(a)) + ) + ) + @test ForwardDiff.derivative(m, 1.0) == 1.0 +end + +@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])) + @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( + 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]) + 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) + + # 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] + # 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 ArgumentError(msg) ReverseDiff.gradient(f, [2.0, 3.0]) +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 "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 "`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() + expected = wrap(d, e, A) + actual = map(value, wrap(track(d, tp), track(e, tp), track(A, tp))) + + @test actual == expected + @test isempty(take_recorded!(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 + a = [1.0, 2.0, 3.0] + c = (4.0, 5.0, 6.0) + + tp = InstructionTape() + y = track(copy(a), tp) .* c + instrs = take_recorded!(tp) + @test y isa TrackedArray + @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) + @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(take_recorded!(tp)) == 1 + + # an abstract element type is checked element by element, skipping non-numbers + 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 + 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))) + @test ReverseDiff.gradient(x -> sum(f(x)), a) ≈ expected + 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] + + @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 + @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) + + # 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 + 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 "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) + @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 + @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 + @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) + @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 + @test ReverseDiff.gradient(t -> sum(t[1] .+ w), v) == [3.0, 0.0, 0.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) + @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) + @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 + # materializing `1/y` would give `Inf` before it ever meets the seed + @test ReverseDiff.gradient(y -> 1.0e-10 * sum(y ./ 1.0e-310), [1.0, 2.0]) == + fill(1.0e-10 / 1.0e-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 + 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 7905e3a5..54486273 100644 --- a/test/derivatives/ElementwiseTests.jl +++ b/test/derivatives/ElementwiseTests.jl @@ -53,20 +53,25 @@ function test_elementwise(f, fopt, x, tp) # record yt = broadcast(fopt, xt) @test yt == y + # 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) == 1 + @test length(tp) == (tracked ? 1 : 0) # reverse out = similar(y, (length(x), length(x))) 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 @@ -472,4 +477,76 @@ 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 + +@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 + +@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 + +@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 diff --git a/test/runtests.jl b/test/runtests.jl index 4f595706..abc1e2ca 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,92 +1,96 @@ 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 + + @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")) + 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