Using ForwardDiff.jl for a function of many variables and parameters Julia

Viewed 1098

The github repo for ForwardDiff.jl has some examples. I am trying to extend the example to take in addition to a vector of variables, a parameter. I cannot get it to work.

This is the example (it is short so I will show it rather than linking)

using ForwardDiff
x = rand(5)
f(x::Vector) = sum(sin, x) .+ prod(tan, x) * sum(sqrt, x);
g = x -> ForwardDiff.gradient(f, x);
g(x) # this outputs the gradient.

I want to modify this since I use functions with multiple parameters as well as variables. As a simple modification I have tried adding a single parameter.

f(x::Vector, y) = (sum(sin, x) .+ prod(tan, x) * sum(sqrt, x)) * y;

I have tried the following to no avail:

fp = x -> ForwardDiff.gradient(f, x); 
fp = x -> ForwardDiff.gradient(f, x, y); 

y = 1
println("test grad: ", fp(x, y))

I get the following error message:

ERROR: LoadError: MethodError: no method matching (::var"#73#74")(::Array{Float64,1}, ::Int64)

A similar question was not answered in 2017. A comment led me to here and it seems the function can only accept one input?

The target function must be unary (i.e., only accept a single argument). ForwardDiff.jacobian is an exception to this rule.

Has this changed? It seems very limited to only be able to differentiate unary functions.

A possible workaround would be to concatenate the list of variables and parameters and then just slice the returned gradient to not include the gradients with respect to the parameters, but this seems silly.

2 Answers

I personally think it makes sense to have this unary-only syntax for ForwardDiff. In your case, you could just pack/unpack x and y into a single vector (nominally x2 below):

julia> using ForwardDiff

julia> x = rand(5)
5-element Array{Float64,1}:
 0.4304735670747184
 0.3939269364431113
 0.7912705403776603
 0.8942024934250143
 0.5724373306715196

julia> f(x::Vector, y) = (sum(sin, x) .+ prod(tan, x) * sum(sqrt, x)) * y;

julia> y = 1
1

julia> f(x2::Vector) = f(x2[1:end-1], x2[end]) % unpacking in f call
f (generic function with 2 methods)

julia> fp = x -> ForwardDiff.gradient(f, x);

julia> println("test grad: ", fp([x; y])) % packing in fp call
test grad: [2.6105844240785796, 2.741442601659502, 1.9913192377198885, 1.9382805843854594, 2.26202717745402, 3.434350946190029]

But my preference would be to explicitly name the partial derivatives differently:

julia> ∂f∂x(x,y) = ForwardDiff.gradient(x -> f(x,y), x)
∂f∂x (generic function with 1 method)

julia> ∂f∂y(x,y) = ForwardDiff.derivative(y -> f(x,y), y)
∂f∂y (generic function with 1 method)

julia> ∂f∂x(x, y)
5-element Array{Float64,1}:
 2.6105844240785796
 2.741442601659502
 1.9913192377198885
 1.9382805843854594
 2.26202717745402

julia> ∂f∂y(x, y)
3.434350946190029

Here's a quick attempt at a function which takes multiple arguments, the same signature as Zygote.gradient:

julia> using ForwardDiff, Zygote

julia> multigrad(f, xs...) = ntuple(length(xs)) do i
         g(y) = f(ntuple(j -> j==i ? y : xs[j], length(xs))...)
         xs[i] isa AbstractArray ? ForwardDiff.gradient(g, xs[i]) : 
           xs[i] isa Number ? ForwardDiff.derivative(g, xs[i]) : nothing
         end;

julia> f1(x,y,z) = sum(x.^2)/y;

julia> multigrad(f1, [1,2,3], 4)
([0.5, 1.0, 1.5], -0.875)

julia> Zygote.gradient(f1, [1,2,3], 4)
([0.5, 1.0, 1.5], -0.875)

For a function with several scalar arguments, this evaluates each derivative separately, and perhaps it would be more efficient to use one evaluation with some Dual(x, (dx, dy, dz)). With large-enough array arguments, ForwardDiff.gradient will already perform multiple evaluations, each with some number of perturbations (the chunk size, which you can control).

Related