AD with respect to Flux parameters

Specific Domains Machine Learning
1

I am trying to compute gradient of a function with respect to Flux parameters.
Example:

using Flux,  ForwardDiff

function f(x::Vector)
    return [x[2]; -x[1] + (1 - x[1]^2) * x[2]]
end

# Function approximation
n = 2; 
fhat = Chain(Dense(2,n),Dense(n,n,sigmoid),Dense(n,1));
ps = Flux.params(fhat);

# Define vector field and divergence
F(x) = f(x) * uhat(x)[1];    
divF(x) = Flux.tr(ForwardDiff.jacobian(F,x)); 

# Compute gradients w.r.t function parameters
x = [1,2];
grads1 = gradient(() -> sum(F(x)), ps); @show grads[ps[1]] # Works as expected.
grads2 = gradient(() -> divF(x), ps);   @show grads[ps[1]] # Nothing.

grads2 returns nothing
How can I obtain derivative of divF(x) w.r.t ps?

Thanks.

1 Like
2

ForwardDiff.jacobian(F, x) does not know about derivatives with respect to F. Maybe we should make that an error again…

(Your variable names require some guessing.)

1 Like