Support for ReverseDiff

[amrods]

It'd be great to add support for ReverseDiff. ForwardDiff is great, and even FiniteDiff works fine. But with ReverseDiff I get an error:

using ComponentArrays
using ForwardDiff
using ReverseDiff
using FiniteDiff
using Parameters: @unpack

ca = ComponentArray(x = [1, 2], θ = [1.0, 100.0], deg = 2)

function F(ca)
    @unpack x, θ, deg = ca
    (θ[1] - x[1])^deg + θ[2] * (x[2] - x[1]^deg)^deg
end

gradsF = ForwardDiff.gradient(ca -> F(ca), ca).x
gradsFD = FiniteDiff.finite_difference_gradient(ca -> F(ca), ca).x
gradsR = ReverseDiff.gradient(ca -> F(ca), ca).x
# ERROR: type TrackedArray has no field x

[jonniedie]

Zygote also works. But yeah, ReverseDiff is one I should also get working because it’s still used in a few places. I’m not sure if ReverseDiff uses ChainRules, but I’ve also been meaning to add that so I can cover more bases. I’ll try to get to this tonight.

[jonniedie]

This is looking a little more difficult than I was expecting it to be. Supposedly there are plans to migrate ReverseDiff to use ChainRules, so I think it would probably be best to wait for that. In the meantime, doing a custom gradient with Zygote like gradsZ = Zygote.gradient(F, ca)[1].x is probably the way to go for things like Optim.

[bgroenks96]

What's the status of this? Does it work now? Last time I checked, ReverseDiff seemed to support ChainRules, though admittedly I'm still a bit hazy on how that all works.

[jonniedie]

Well the good news is I updated my autodiff stuff to use ChainRules instead of Zygote directly. The bad news is ReverseDiff still doesn't use ChainRules. I'm looking into what it's going to take to make this work, but the ReverseDiff docs don't give me a ton to work with for figuring out how to make pullback rules. The only thing I can think of is overloading each derivative function (gradient, gradient!, jacobian, jacobian!, ...) for getproperty and getindex of a ReverseDiff.TrackedArray{V,D,N,<:ComponentArray,DA} where {V,D,N,DA}. If that's the case, I may need some help in the form of a PR or it will be a while until I can get to it.

[bgroenks96]

See #67

[bgroenks96]

Yeah so I kind of need this, possibly soon. What can I do to help?

[bgroenks96]

Although I actually don't necessarily need my parameter array to be a ComponentArray. I guess I should check if non-tracked arrays work with ComponentArrays.

[jonniedie]

It's looking like we might be able to define getindex for a tracked ComponentArray and a Symbol or Val taking some inspiration from here

using ComponentArrays, ReverseDiff
using ReverseDiff: TrackedArray, value, deriv, tape, SpecialInstruction, record! 

const TrackedComponentArray{V, D, N, DA, N, A, Ax} = TrackedArray{V,D,N,ComponentArray{V,N,A,Ax},DA}

function Base.getindex(t::TrackedComponentArray, ind::Union{Symbol, Val})
    tp = tape(t)
    out = TrackedArray(value(t)[ind], deriv(t)[ind], tp)
    record!(tp, SpecialInstruction, (getindex, Val(:generic)), (t, ind), out)
    return out
end

This is enough to get the tracking to work through the getindex operation (getproperty is probably not a good idea because TrackedArray needs to be able to access its fields with that), but actually taking derivatives, gradients, jacobians, etc. requires reaching into the special_reverse_exec! code here. Or maybe just making a separate dispatch for something other than Val(:generic) and patch in from there. I haven't had luck getting that working yet, though.

[bgroenks96]

So without getproperty we would only be able to do carray[:var] right?

[bgroenks96]

I guess I don't really know what special_reverse_exec is doing. What is SpecialInstruction?

[jonniedie]

@bgroenks96 So I guess I didn't need the special instruction thing after all. I tried this out and it seems to work for me, but I've just tested it on pretty simple stuff so far. Try adding this and running your code:

using ComponentArrays
using ReverseDiff: ReverseDiff, TrackedArray, TrackedReal, value, deriv, tape

const TrackedComponentArray{V, D, N, DA, N, A, Ax} = TrackedArray{V,D,N,ComponentArray{V,N,A,Ax},DA}

maybe_tracked_array(val::AbstractArray, der, t) = TrackedArray(val, der, t)
maybe_tracked_array(val, der, t) = TrackedReal(val, der, t)

function Base.getindex(t::TrackedComponentArray, inds::Union{Symbol, Val}...)
    return maybe_tracked_array(value(t)[inds...], deriv(t)[inds...], tape(t))
end

function Base.getproperty(t::TrackedComponentArray, s::Symbol)
    if s in (:value, :deriv, :tape)
        return getfield(t, s)
    else
        return maybe_tracked_array(getproperty(value(t), s), getproperty(deriv(t), s), tape(t))
    end
end

[jonniedie]

I'll do a little more testing and push an update soon.

[bgroenks96]

Hmm, I'm still getting the Type TrackedArray has no field ... error. Am I supposed to change how I use it?

[bgroenks96]

Wait nevermind, that was an urelated problem. Here's the real error:

MethodError: getindex(::TrackedArray{Float64,Float64,0,Base.ReshapedArray{Float64,0,SubArray{Float64,1,Array{Float64,1},Tuple{UnitRange{Int64}},true},Tuple{}},Base.ReshapedArray{Float64,0,SubArray{Float64,1,Array{Float64,1},Tuple{UnitRange{Int64}},true},Tuple{}}}) is ambiguous. Candidates:
  getindex(t::TrackedArray, i::AbstractRange...) in ReverseDiff at /home/brian/.julia/packages/ReverseDiff/iHmB4/src/tracked.jl:299
  getindex(t::TrackedArray, i::Colon...) in ReverseDiff at /home/brian/.julia/packages/ReverseDiff/iHmB4/src/tracked.jl:299
  getindex(t::TrackedArray, i::Int64...) in ReverseDiff at /home/brian/.julia/packages/ReverseDiff/iHmB4/src/tracked.jl:338
  getindex(t::TrackedArray, i::Union{Colon, AbstractRange}...) in ReverseDiff at /home/brian/.julia/packages/ReverseDiff/iHmB4/src/tracked.jl:299
Possible fix, define
  getindex(::TrackedArray)
isassigned(::TrackedArray{Float64,Float64,0,Base.ReshapedArray{Float64,0,SubArray{Float64,1,Array{Float64,1},Tuple{UnitRange{Int64}},true},Tuple{}},Base.ReshapedArray{Float64,0,SubArray{Float64,1,Array{Float64,1},Tuple{UnitRange{Int64}},true},Tuple{}}}) at abstractarray.jl:408
...

And it seems to occur because somehow my sub-array in ComponentVector becomes a TrackedReal?

TrackedReal{Float64,Float64,TrackedArray{Float64,Float64,1,ComponentVector{Float64,SubArray...},ComponentVector{Float64,SubArray...}}}[TrackedReal<4L7>(0.0, 0.0, 4rg, 1, LLy)]

I don't know, could be an issue in my code as well. Maybe I should just simplify down to a more basic example (like not involving sensitivities).

[bgroenks96]

Also this problem apparently also exists for DEDataArray because getproperty fails when it gets turned into TrackedArray @ChrisRackauckas