Support arrays with non-concrete eltype (fixes #264)

src/derivatives/propagation.jl

@@ -30,8 +30,19 @@ index_bound(x::AbstractArray, ::AbstractArray{T,N}) where {T,N} = CartesianIndex
 
 @inline increment_deriv!(t::TrackedArray, x::AbstractArray, i) = (t.deriv[i] += x[i]; nothing)
 @inline increment_deriv!(t::TrackedArray, x::Real, i) = (t.deriv[i] += x; nothing)
-@inline increment_deriv!(t::AbstractArray, x::AbstractArray, i) = increment_deriv!(t[i], x[i])
-@inline increment_deriv!(t::AbstractArray, x::Real, i) = increment_deriv!(t[i], x)
+
+# `istracked(t)` only promises that an element *may* be tracked
+@inline function increment_deriv!(t::AbstractArray, x::AbstractArray, i)
+    ti = t[i]
+    istracked(ti) && increment_deriv!(ti, x[i])
+    return nothing
+end
+
+@inline function increment_deriv!(t::AbstractArray, x::Real, i)
+    ti = t[i]
+    istracked(ti) && increment_deriv!(ti, x)
+    return nothing
+end
 
 function increment_deriv!(t::AbstractArray, x)
     for i in eachindex(t)
@@ -53,8 +64,19 @@ end
 
 @inline decrement_deriv!(t::TrackedArray, x::AbstractArray, i) = (t.deriv[i] -= x[i]; nothing)
 @inline decrement_deriv!(t::TrackedArray, x::Real, i) = (t.deriv[i] -= x; nothing)
-@inline decrement_deriv!(t::AbstractArray, x::AbstractArray, i) = decrement_deriv!(t[i], x[i])
-@inline decrement_deriv!(t::AbstractArray, x::Real, i) = decrement_deriv!(t[i], x)
+
+# `istracked(t)` only promises that an element *may* be tracked
+@inline function decrement_deriv!(t::AbstractArray, x::AbstractArray, i)
+    ti = t[i]
+    istracked(ti) && decrement_deriv!(ti, x[i])
+    return nothing
+end
+
+@inline function decrement_deriv!(t::AbstractArray, x::Real, i)
+    ti = t[i]
+    istracked(ti) && decrement_deriv!(ti, x)
+    return nothing
+end
 
 function decrement_deriv!(t::AbstractArray, x)
     for i in eachindex(t)

src/macros.jl

@@ -374,11 +374,13 @@ macro grad_from_chainrules(fcall)
 end
 
 _add_to_deriv!(d1, d2) = nothing
-function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray{<:TrackedReal}}, d2::AbstractThunk)
-    increment_deriv!(d1, unthunk(d2))
+function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, d2::AbstractThunk)
+    istracked(d1) && increment_deriv!(d1, unthunk(d2))
+    return nothing
 end
-function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray{<:TrackedReal}}, d2)
-    increment_deriv!(d1, d2)
+function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, d2)
+    istracked(d1) && increment_deriv!(d1, d2)
+    return nothing
 end
 function getargs_expr(args_with_types)
     expr = Expr(:tuple)

src/tracked.jl

@@ -97,7 +97,10 @@ const TrackedVecOrMat{V,D} = Union{TrackedVector{V,D}, TrackedMatrix{V,D}}
 istracked(x) = false
 istracked(::TrackedReal) = true
 istracked(::TrackedArray) = true
-istracked(::AbstractArray{T}) where {T} = T <: TrackedReal || !(isconcretetype(T))
+function istracked(::AbstractArray{T}) where {T}
+    # can any tracked type inhabit `T`?
+    return typeintersect(T, TrackedReal) !== Union{} || typeintersect(T, TrackedArray) !== Union{}
+end
 
 @inline value(x) = x
 @inline value(x::AbstractArray) = istracked(x) ? map(value, x) : x

test/ChainRulesTests.jl

@@ -204,6 +204,50 @@ end
     @test results[1] == fill(38, size(inputs[1])) # 38 = 3 + 5 * 7
 end
 
+### Thunked tangents
+# Lazy tangents have to be unthunked: `increment_deriv!` only accepts arrays and reals.
+
+f_thunk(x) = sum(4x .+ 1)
+function ChainRulesCore.rrule(::typeof(f_thunk), x)
+    # `InplaceableThunk` is what ChainRules' matrix rules hand back
+    pb_f_thunk(Δ) = (NoTangent(), InplaceableThunk(ȳ -> ȳ .+= 3Δ, @thunk(fill(3Δ, size(x)))))
+    return f_thunk(x), pb_f_thunk
+end
+ReverseDiff.@grad_from_chainrules f_thunk(x::ReverseDiff.TrackedArray)
+
+f_thunk(x::Real) = 4x + 1
+function ChainRulesCore.rrule(::typeof(f_thunk), x::Real)
+    pb_f_thunk_scalar(Δ) = (NoTangent(), @thunk(3Δ))
+    return f_thunk(x), pb_f_thunk_scalar
+end
+ReverseDiff.@grad_from_chainrules f_thunk(x::ReverseDiff.TrackedReal)
+
+# `@grad_from_chainrules` dispatches on `TrackedArray`, so only `@grad` reaches an array
+# that merely *holds* tracked elements
+g_thunk(a, b) = sum(a) + sum(b)
+function g_thunk(a::AbstractVector, b::AbstractVector)
+    if ReverseDiff.istracked(a) || ReverseDiff.istracked(b)
+        return ReverseDiff.track(g_thunk, a, b)
+    else
+        return sum(a) + sum(b)
+    end
+end
+ReverseDiff.@grad function g_thunk(a, b)
+    av, bv = ReverseDiff.value(a), ReverseDiff.value(b)
+    pb_g_thunk(Δ) = (@thunk(fill(3Δ, size(av))), @thunk(fill(5Δ, size(bv))))
+    return sum(av) + sum(bv), pb_g_thunk
+end
+
+@testset "thunked tangents" begin
+    x = rand(3, 3)
+    @test ReverseDiff.gradient(f_thunk, x) == fill(3, size(x))
+    @test ReverseDiff.gradient(u -> f_thunk(u[1]), [2.0]) == [3.0]
+
+    # the thunk for the untracked argument is dropped rather than forced
+    @test ReverseDiff.gradient(u -> g_thunk(Real[u[1], u[2]], [1.0, 2.0]), [1.0, 2.0]) == [3.0, 3.0]
+    @test ReverseDiff.gradient(u -> g_thunk([1.0, 2.0], Real[u[1], u[1]]), [1.0, 2.0]) == [10.0, 0.0]
+end
+
 ### Isolated Scope
 module IsolatedModuleForTestingScoping
 using ChainRulesCore

test/TrackedTests.jl

@@ -98,6 +98,24 @@ t = TrackedArray(v, d, tp)
 @test ReverseDiff.istracked(Any[1])
 @test ReverseDiff.istracked([TrackedReal(1, 1)])
 
+# eltypes that cannot hold a tracked value (`TrackedReal <: Real`, but not `<: AbstractFloat`)
+@test !(ReverseDiff.istracked(AbstractString["a"]))
+@test !(ReverseDiff.istracked(Integer[1]))
+@test !(ReverseDiff.istracked(AbstractFloat[1.0]))
+# eltypes that can
+@test ReverseDiff.istracked(Real[1.0])
+@test ReverseDiff.istracked(Number[1.0])
+@test ReverseDiff.istracked(AbstractVector[[1.0]])
+@test ReverseDiff.istracked(Union{Float64,TrackedReal{Float64,Float64,Nothing}}[1.0])
+
+# a vector of `TrackedArray`s is tracked, and `value` and `tape` see through it
+let ta = TrackedArray(rand(2), rand(2), InstructionTape()), v = [ta, ta]
+    @test ReverseDiff.istracked(v)
+    @test ReverseDiff.value(v) == [ReverseDiff.value(ta), ReverseDiff.value(ta)]
+    @test eltype(ReverseDiff.value(v)) === Vector{Float64}
+    @test ReverseDiff.tape(v) === ReverseDiff.tape(ta)
+end
+
 # value #
 #-------#
 

test/api/GradientTests.jl

@@ -366,4 +366,44 @@ g_untracked(x, y) = 2.0
     @test_throws MethodError ReverseDiff.gradient!(DiffResults.GradientResult(x), g251, (x, y))
 end
 
+############################################################################################
+
+@testset "arrays with non-concrete eltype (#264)" begin
+    # the MWE from the issue
+    f(u) = last((Real[1.0, 2.0] * u[], u[]))
+    @test ReverseDiff.gradient(f, [2.0]) == [1.0]
+
+    # each operator reaches a different propagation path
+    @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] .* u[]), [2.0]) == [3.0]
+    @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] .+ u[]), [2.0]) == [2.0]
+    @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] .- u[]), [2.0]) == [-2.0]
+    @test ReverseDiff.gradient(u -> sum(u[] .- Real[1.0, 2.0]), [2.0]) == [2.0]
+    @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] ./ u[]), [2.0]) == [-0.75]
+    @test ReverseDiff.gradient(u -> sum(u[] ./ Real[1.0, 2.0]), [2.0]) == [1.5]
+
+    # tracked elements must still accumulate, including when one input repeats
+    @test ReverseDiff.gradient(u -> sum(Real[u[1], u[2]] .* 2.0), [2.0, 3.0]) == [2.0, 2.0]
+    @test ReverseDiff.gradient(u -> sum(Real[u[1], u[1]] .* 2.0), [2.0, 3.0]) == [4.0, 0.0]
+
+    # mixed tracked/untracked elements, and `Any` eltype
+    @test ReverseDiff.gradient(u -> sum(Real[3.0, u[2]] .* 2.0), [2.0, 3.0]) == [0.0, 2.0]
+    @test ReverseDiff.gradient(u -> sum(2.0 .- Real[3.0, u[2]]), [2.0, 3.0]) == [0.0, -1.0]
+    @test ReverseDiff.gradient(u -> sum(Any[3.0, u[2]] .* 2.0), [2.0, 3.0]) == [0.0, 2.0]
+
+    # linalg, the other APIs, and tape replay
+    A = Real[1.0 2.0; 3.0 4.0]
+    @test ReverseDiff.gradient(u -> sum(A * u), [1.0, 2.0]) == [4.0, 6.0]
+    @test ReverseDiff.jacobian(u -> A * u, [1.0, 2.0]) == [1.0 2.0; 3.0 4.0]
+    @test ReverseDiff.hessian(u -> sum((A * u) .^ 2), [1.0, 2.0]) == 2 * (A' * A)
+
+    tape = ReverseDiff.GradientTape(u -> sum(A * u), [1.0, 2.0])
+    @test ReverseDiff.gradient!(tape, [5.0, 6.0]) == [4.0, 6.0]
+
+    # `@grad` rules propagate to non-concrete eltype inputs (#145)
+    @test ReverseDiff.gradient(x -> sum([Real[x[1], x[2]]; x[3]]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 1.0]
+    @test ReverseDiff.gradient(x -> sum([Any[x[1], x[2]]; x[3]]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 1.0]
+    @test ReverseDiff.gradient(x -> sum([Real[5.0, x[2]]; x[3]]), [1.0, 2.0, 3.0]) == [0.0, 1.0, 1.0]
+    @test ReverseDiff.gradient(x -> sum([Real[x[1], x[2]] x[3]*ones(2)]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 2.0]
+end
+
 end # module