Using automatic differentiation on a function that makes use of a preallocated array in Julia

Question

My long subject title pretty much covers it.

I have managed to isolate my much bigger problem in the following contrived example below. I cannot figure out where the problem exactly is, though I imagine it has something to do with the type of the preallocated array?

using ForwardDiff

function test()

    A = zeros(1_000_000)

    function objective(A, value)
        for i=1:1_000_000
            A[i] = value[1]
        end

        return sum(A)
    end

    helper_objective = v -> objective(A, v)

    ForwardDiff.gradient(helper_objective, [1.0])

end

The error reads as follows:

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{getfield(Main, Symbol("##69#71")){Array{Float64,1},getfield(Main, Symbol("#objective#70")){Array{Float64,1}}},Float64},Float64,1})

In my own problem (not described here) I have a function that I need to optimise using Optim, and the automatic differentiation it offers, and this function makes use of a big matrix that I would like to preallocate in order to speed up my code. Many thanks.

Answer 1

If you look at http://www.juliadiff.org/ForwardDiff.jl/latest/user/limitations.html you find:

The target function must be written generically enough to accept numbers of type T<:Real as input (or arrays of these numbers) (...) This also means that any storage assigned used within the function must be generic as well.

with the example here https://github.com/JuliaDiff/ForwardDiff.jl/issues/136#issuecomment-237941790.

This means that you could do something like this:

function test()
    function objective(value)
        for i=1:1_000_000
            A[i] = value[1]
        end
        return sum(A)
    end
    A = zeros(ForwardDiff.Dual{ForwardDiff.Tag{typeof(objective), Float64},Float64,1}, 1_000_000)
    ForwardDiff.gradient(objective, [1.0])
end

But I would not assume that this will save you much allocations as it is type unstable.

What you can do is wrap objective and A in a module like this:

using ForwardDiff

module Obj

using ForwardDiff

function objective(value)
    for i=1:1_000_000
        A[i] = value[1]
    end
    return sum(A)
end
const A = zeros(ForwardDiff.Dual{ForwardDiff.Tag{typeof(objective), Float64},Float64,1}, 1_000_000)

end

And now this:

ForwardDiff.gradient(Obj.objective, [1.0])

should be fast.

EDIT

Also this works (although it is type unstable but in a less problematic place):

function test()::Vector{Float64}
    function objective(A, value)
        for i=1:1_000_000
            A[i] = value[1]
        end

        return sum(A)
    end
    helper_objective = v -> objective(A, v)
    A = Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(helper_objective), Float64},Float64,1}}(undef, 1_000_000)
    ForwardDiff.gradient(helper_objective, [1.0])
end