Put broadcasting on Julia's documented custom-broadcast interface

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"

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

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

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! #
 ##############################

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)

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

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