ForwardDiff and Array of Arrays

Is it possible to get ForwardDiff to work with arrays of arrays? I’m thinking about something like this:

using LinearAlgebra
using ForwardDiff
using Random

function f(X)
   return X[1]⋅cross(X[2],X[3]) 
end

Random.seed!(100)
X = [randn(3) for j=1:3]

If I try to use automatic differentiation on this:

∇f = X->ForwardDiff.gradient(f,X)
∇f(X)

I get the error:

MethodError: no method matching one(::Type{Array{Float64,1}})
Closest candidates are:
  one(!Matched::Type{Missing}) at missing.jl:103
  one(!Matched::BitArray{2}) at bitarray.jl:400
  one(!Matched::Missing) at missing.jl:100
  ...

Stacktrace:
 [1] macro expansion at /Users/guardian/.julia/packages/ForwardDiff/Asf4O/src/partials.jl:0 [inlined]
 [2] single_seed(::Type{ForwardDiff.Partials{3,Array{Float64,1}}}, ::Val{1}) at /Users/guardian/.julia/packages/ForwardDiff/Asf4O/src/partials.jl:10
 [3] macro expansion at /Users/guardian/.julia/packages/ForwardDiff/Asf4O/src/apiutils.jl:0 [inlined]
 [4] construct_seeds(::Type{ForwardDiff.Partials{3,Array{Float64,1}}}) at /Users/guardian/.julia/packages/ForwardDiff/Asf4O/src/apiutils.jl:53
 [5] ForwardDiff.GradientConfig(::typeof(f), ::Array{Array{Float64,1},1}, ::ForwardDiff.Chunk{3}, ::ForwardDiff.Tag{typeof(f),Array{Float64,1}}) at /Users/guardian/.julia/packages/ForwardDiff/Asf4O/src/config.jl:121 (repeats 3 times)
 [6] gradient(::Function, ::Array{Array{Float64,1},1}) at /Users/guardian/.julia/packages/ForwardDiff/Asf4O/src/gradient.jl:15
 [7] (::var"#11#12")(::Array{Array{Float64,1},1}) at ./In[15]:1
 [8] top-level scope at In[16]:1

Is this just asking too much of ForwardDiff?

I don’t know about this particular use case but I have used ForwardDiff with an array of tables in the following manner

pmap(x -> ForwardDiff.gradient!(Array{Float64}(undef, 1, J), y -> loglike(y, x), beta), data) 

where data is an array of tables an beta is a 1 x J row-vector. The output returned was an array of array of floats. Maybe you can adapt something similar here? If you don’t want to use pmap, try mapslices

maybe using a vector instead of a vector of vectors and reshaping?

function myf(x)
    xvec=reshape(x,3,3)
    return xvec[:,1]⋅cross(xvec[:,2],xvec[:,3]) 
 end

 ∇myf = X->ForwardDiff.gradient(myf,X)
 x0 = randn(9)
 dx = reshape(∇myf(x0),3,3) #here, dx[:,1] ==  X[1]

if you really (emphasis on really) need an array of arrays, just add this code and ∇f should work with your input, returning an array of arrays of the same form as the input.

function _f(x)
    xvec=reshape(x,3,3)
    return xvec[:,1]⋅cross(xvec[:,2],xvec[:,3]) 
 end
function ∇f(X)
    X2 = reduce((x,y)->cat(x,y,dims=1),X)
    ∇X = ForwardDiff.gradient(myf,X2)
    return [∇X[:,j] for j=1:3]
end

why is this necessary? ForwardDiff accepts functions with only 2 types of input: a real number, and a vector of real numbers. any other structure must be transformed to any of those cases to be differentiated via ForwardDiff.

1 Like

Like @longemen3000 writes, ForwardDiff.jl needs so see a plain Array. But you can use ArraysOfArrays.jl to reshape things in comfort, and you probably want to use StaticArrays.jl for your inner vectors, since the cross product requires a fixed length anyhow.

Leaving your f(x) unchanged:

using LinearAlgebra
using ForwardDiff
using Random

function f(X)
   return X[1]⋅cross(X[2],X[3]) 
end

Random.seed!(100)

you can do

using ArraysOfArrays, StaticArrays
vv3(X::Matrix) = nestedview(X, SVector{3})

X = rand(3, 3)

∇f = X -> ForwardDiff.gradient(X -> f(vv3(X)),X)
vv3(∇f(X))

vv3(X) will reshape a matrix into a vector of static vectors, but ForwardDiff will see the flat matrix.

Thanks for the suggestions. Eventually, I was able to use Zygote.jl. This was fine for me, as I actually just wanted to verify that a hand coded gradient was correct.

2 Likes