Divide instead of multiplying by cached reciprocals (fixes #239)

Project.toml

@@ -1,6 +1,6 @@
 name = "ReverseDiff"
 uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-version = "1.18.0"
+version = "1.18.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

ext/StatisticsExt.jl

@@ -14,7 +14,7 @@ end
 @noinline function ReverseDiff.special_reverse_exec!(instruction::SpecialInstruction{typeof(mean)})
     input = instruction.input
     output = instruction.output
-    istracked(input) && increment_deriv!(input, inv(length(input)) * deriv(output))
+    istracked(input) && increment_deriv!(input, deriv(output) / length(input))
     unseed!(output)
     return nothing
 end

src/derivatives/elementwise.jl

@@ -539,39 +539,28 @@ end
 # ./ #
 #----#
 
-numer_partials(d::Real) = Ref(inv(d))
-numer_partials(d::AbstractArray) = broadcast(inv, d)
-numer_partials!(out::Ref, d) = (out[] = inv(d); nothing)
-numer_partials!(out::AbstractArray, d) = (broadcast!(inv, out, d); nothing)
-
-denom_partials_kernel(n::Real, d::Real) =  -(n / (d * d))
-denom_partials(n::Real, d::Real) = Ref(denom_partials_kernel(n, d))
+denom_partials_kernel(n::Real, d::Real) = -(n / d) / d
 denom_partials(n, d) = broadcast(denom_partials_kernel, n, d)
-denom_partials!(out::Ref, n, d) = (out[] = denom_partials_kernel(n, d); nothing)
 denom_partials!(out::AbstractArray, n, d) = (broadcast!(denom_partials_kernel, out, n, d); nothing)
 
-rdiv_cache(x, y) = (numer_partials(value(y)), denom_partials(value(x), value(y)))
-
 function broadcast_rdiv(x, y, ::Type{D}) where D
     tp = tape(x, y)
     out = track(value(x) ./ value(y), D, tp)
-    n_partials, d_partials = rdiv_cache(x, y)
-    cache = (n_partials, d_partials,
-             index_bound(x, out), index_bound(y, out),
-             index_bound(n_partials, out), index_bound(d_partials, out))
+    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)
-    n_partials, d_partials = instruction.cache
+    d_partials = first(instruction.cache)
     output = instruction.output
     pull_value!(a)
     pull_value!(b)
     broadcast!(/, value(output), a_value, b_value)
-    istracked(a) && numer_partials!(n_partials, b_value)
     istracked(b) && denom_partials!(d_partials, a_value, b_value)
     return nothing
 end
@@ -580,9 +569,8 @@ end
     a, b = instruction.input
     output = instruction.output
     output_deriv = deriv(output)
-    n_partials, d_partials, a_bound, b_bound,
-    n_partials_bound, d_partials_bound = instruction.cache
-    istracked(a) && broadcast_increment_deriv!(a, output_deriv, n_partials, a_bound, n_partials_bound)
+    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
@@ -594,23 +582,21 @@ end
 function broadcast_ldiv(x, y, ::Type{D}) where D
     tp = tape(x, y)
     out = track(value(x) .\ value(y), D, tp)
-    n_partials, d_partials = rdiv_cache(y, x)
-    cache = (n_partials, d_partials,
-             index_bound(x, out), index_bound(y, out),
-             index_bound(n_partials, out), index_bound(d_partials, out))
+    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)
-    n_partials, d_partials = instruction.cache
+    d_partials = first(instruction.cache)
     output = instruction.output
     pull_value!(a)
     pull_value!(b)
     broadcast!(\, value(output), a_value, b_value)
-    istracked(b) && numer_partials!(n_partials, a_value)
     istracked(a) && denom_partials!(d_partials, b_value, a_value)
     return nothing
 end
@@ -619,10 +605,9 @@ end
     a, b = instruction.input
     output = instruction.output
     output_deriv = deriv(output)
-    n_partials, d_partials, a_bound, b_bound,
-    n_partials_bound, d_partials_bound = instruction.cache
+    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_deriv!(b, output_deriv, n_partials, b_bound, n_partials_bound)
+    istracked(b) && broadcast_increment_div_deriv!(b, output_deriv, value(a), b_bound, a_bound)
     unseed!(output)
     return nothing
 end
@@ -631,15 +616,11 @@ end
 #----#
 
 base_partials_kernel(b::Real, e::Real) = e * b^(e - 1)
-base_partials(b::Real, e::Real) = Ref(base_partials_kernel(b, e))
 base_partials(b, e) = broadcast(base_partials_kernel, b, e)
-base_partials!(out::Ref, b, e) = (out[] = base_partials_kernel(b, e); nothing)
 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::Real, e::Real) = Ref(exp_partials_kernel(b, e))
 exp_partials(b, e) = broadcast(exp_partials_kernel, b, e)
-exp_partials!(out::Ref, b, e) = (out[] = exp_partials_kernel(b, e); nothing)
 exp_partials!(out::AbstractArray, b, e) = (broadcast!(exp_partials_kernel, out, b, e); nothing)
 
 function pow_cache(x, y)

src/derivatives/propagation.jl

@@ -7,10 +7,8 @@ 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.
-These special cases should be decently fast, but some might be prone to numerical error if
-the accumulated derivative becomes too large compared to the individual terms being added to
-it. This can be overcome by using the divide-and-conquer strategy from Base.mapreducedim,
-but that strategy is less cache efficient and more complicated to implement.
+Accumulations into a scalar derivative go through `sum`, whose pairwise reduction is both more
+accurate and faster than a sequential loop.
 
 A lot of the code here is pretty repetitive, and only covers the use cases that actually
 arise from the derivative definitions implemented elsewhere in ReverseDiff. At some point,
@@ -81,28 +79,28 @@ end
 
 function diffresult_increment_deriv!(input::AbstractArray, x::AbstractArray,
                                      results, p::Int)
-    for i in eachindex(x)
+    for i in eachindex(x, results)
         increment_deriv!(input, x[i] * getpartial(results[i], p), i)
     end
     return nothing
 end
 
 function diffresult_increment_deriv!(input::AbstractArray, x::AbstractArray,
                                      results, p::Int, bound::CartesianIndex)
-    for i in CartesianIndices(size(x))
-        increment_deriv!(input, x[i] * getpartial(results[i], p), min(bound, i))
+    axes(x) == axes(results) ||
+        throw(DimensionMismatch("`x` and `results` must have the same indices"))
+    for (xi, r, i) in zip(x, results, CartesianIndices(size(x)))
+        increment_deriv!(input, xi * getpartial(r, p), min(bound, i))
     end
     return nothing
 end
 
 function diffresult_increment_deriv!(input::TrackedReal, x::AbstractArray,
                                      results, p::Int, ::Nothing)
+    inds = eachindex(x, results)
+    isempty(inds) && return nothing
     pull_deriv!(input)
-    input_deriv = input.deriv
-    for i in eachindex(results)
-        input_deriv += x[i] * getpartial(results[i], p)
-    end
-    input.deriv = input_deriv
+    input.deriv += sum(i -> x[i] * getpartial(results[i], p), inds)
     push_deriv!(input)
     return nothing
 end
@@ -116,74 +114,63 @@ end
 
 function broadcast_increment_deriv!(input::AbstractArray, x::AbstractArray,
                                     bound::CartesianIndex)
-    for i in CartesianIndices(size(x))
-        increment_deriv!(input, x[i], min(bound, i))
+    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 = input.deriv
-    for i in x
-        input_deriv += i
-    end
-    input.deriv = input_deriv
+    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!(input::AbstractArray, x::AbstractArray,
-                                    partials::AbstractArray,
-                                    input_bound::CartesianIndex,
-                                    partials_bound::CartesianIndex)
-    for i in CartesianIndices(size(x))
-        current_deriv = x[i] * partials[min(partials_bound, i)]
+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))
     end
     return nothing
 end
 
-function broadcast_increment_deriv!(input::TrackedReal, x::AbstractArray,
-                                    partials::AbstractArray, ::Nothing, ::CartesianIndex)
+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 = input.deriv
-    for i in eachindex(x)
-        input_deriv += x[i] * partials[i]
-    end
-    input.deriv = input_deriv
+    input.deriv += sum(i -> op(x[i], partials[i]), inds)
     push_deriv!(input)
     return nothing
 end
 
 # with partial scalar #
 #---------------------#
 
-@inline function broadcast_increment_deriv!(input, x, partial::Ref, i, j)
-    return broadcast_increment_deriv!(input, x, partial[], i, j)
-end
-
-function broadcast_increment_deriv!(input::AbstractArray, x::AbstractArray,
-                                    partial::Real, input_bound::CartesianIndex,
-                                    ::Nothing)
-    for i in CartesianIndices(size(x))
-        increment_deriv!(input, x[i] * partial, min(input_bound, i))
-    end
-    return nothing
-end
-
-function broadcast_increment_deriv!(input::TrackedReal, x::AbstractArray,
-                                    partial::Real, ::Nothing, ::Nothing)
-    pull_deriv!(input)
-    input_deriv = input.deriv
-    for i in eachindex(x)
-        input_deriv += x[i] * partial
+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))
     end
-    input.deriv = input_deriv
-    push_deriv!(input)
     return nothing
 end
 
@@ -193,19 +180,16 @@ end
 
 function broadcast_decrement_deriv!(input::AbstractArray, x::AbstractArray,
                                     bound::CartesianIndex)
-    for i in CartesianIndices(size(x))
-        decrement_deriv!(input, x[i], min(bound, i))
+    for (xi, i) in zip(x, CartesianIndices(size(x)))
+        decrement_deriv!(input, xi, min(bound, i))
     end
     return nothing
 end
 
 function broadcast_decrement_deriv!(input::TrackedReal, x::AbstractArray, ::Nothing)
+    isempty(x) && return nothing
     pull_deriv!(input)
-    input_deriv = input.deriv
-    for i in x
-        input_deriv -= i
-    end
-    input.deriv = input_deriv
+    input.deriv -= sum(x)
     push_deriv!(input)
     return nothing
 end

test/PrecisionTests.jl

@@ -0,0 +1,56 @@
+module PrecisionTests
+
+using ReverseDiff, Test, LinearAlgebra, Statistics
+
+include("utils.jl")
+
+# gradients must be accurate to the working precision (#239)
+setprecision(BigFloat, 256) do
+    x = big.(rand(5)) ./ 3
+    n = length(x)
+    T = eltype(x)
+    rtol = 8 * eps(T)
+
+    for (f, ∇f) in ((mean,                    y -> fill(one(T) / n, n)),
+                    (y -> mean(y .^ 3),       y -> 3 .* y .^ 2 ./ n),
+                    (y -> sum(y ./ 3),        y -> fill(one(T) / 3, n)),
+                    (y -> sum(y ./ 3.0),      y -> fill(one(T) / 3, n)),
+                    (y -> sum(3 .\ y),        y -> fill(one(T) / 3, n)),
+                    (y -> sum(3.0 .\ y),      y -> fill(one(T) / 3, n)),
+                    (sum,                     y -> fill(one(T), n)),
+                    (prod,                    y -> prod(y) ./ y),
+                    (norm,                    y -> y ./ norm(y)),
+                    (var,                     y -> 2 .* (y .- mean(y)) ./ (n - 1)),
+                    (std,                     y -> (y .- mean(y)) ./ ((n - 1) * std(y))),
+                    (y -> dot(y, y),          y -> 2 .* y),
+                    (y -> sum(3 ./ y),        y -> -3 ./ y .^ 2),
+                    (y -> sum(y .\ 3),        y -> -3 ./ y .^ 2),
+                    (y -> sum(y ./ (y .+ 1)), y -> 1 ./ (y .+ 1) .^ 2),
+                    (y -> sum(exp.(y)),       y -> exp.(y)))
+        test_println("BigFloat gradients", f)
+        @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
+    end
+end
+
+# a subnormal denominator must not overflow the gradient
+let f = y -> 1e-10 * sum(y ./ 1e-310)
+    test_println("Float64 gradients", f)
+    @test ReverseDiff.gradient(f, [1.0, 2.0]) == fill(1e-10 / 1e-310, 2)
+end
+
+# the denominator adjoint is `-n / d^2`, which must not overflow while it is still finite
+for (n, d) in ((1e-200, 1e-200), (1e200, 1e200))
+    ∂ = Float64(-big(n) / big(d)^2)
+    for f in (y -> sum([n] ./ y), y -> sum(y .\ [n]))
+        test_println("Float64 gradients", f)
+        @test ReverseDiff.gradient(f, [d]) ≈ [∂] rtol=8*eps(Float64)
+    end
+end
+
+end # module

test/runtests.jl

@@ -42,6 +42,12 @@ test_println(kind, f, pad="  ") = println(pad, "testing $(kind): `$(f)`...")
         println("done (took $t seconds).")
     end
 
+    @testset "PrecisionTests" begin
+        println("running PrecisionTests...")
+        t = @elapsed include(joinpath(TESTDIR, "PrecisionTests.jl"))
+        println("done (took $t seconds).")
+    end
+
     @testset "ElementWiseTests" begin
         println("running ElementwiseTests...")
         t = @elapsed include(joinpath(TESTDIR, "derivatives/ElementwiseTests.jl"))