Jacobian and Hessian

[MikeInnes]

Should be easy to port from Tracker's versions.

[baggepinnen]

The current hessian function gives stack overflow on the following

using Flux
inp     = randn(30)      
f = Chain(Dense(30, 30, tanh),Dense(30, 30, tanh),Dense(30, 30, tanh),Dense(30, 1, tanh))
Zygote.hessian(f, inp)

[gdkrmr]

I need a jacobian, this works:

function jacobian(f, x)
    y = f(x)
    n = length(y)
    m = length(x)
    T = eltype(y)
    j = Array{T, 2}(undef, n, m)
    for i in 1:n
        j[i, :] .= gradient(x -> f(x)[i], x)[1]
    end
    return j
end

I am sure this can still be optimized but I don't know enough about the internals.

[ferrolho]

I am trying to get the Jacobian of a R^n -> R^m function, where n > 1 and m > 1 (vector input and vector output); for example, a function like func_1 below.

function func_1(x::Array{T}) where T  # out-of-place
    dx = zeros(T, 3)
    dx[1] = 2x[1] + √x[3]
    dx[2] = 3x[1] + sin(x[4])
    dx[3] = 4x[1] + cos(x[5])
    dx
end

When I try to get the Jacobian with the function posted above by @gdkrmr with jacobian(func_1, rand(5)) I get the following error:

Mutating arrays is not supported

Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] (::getfield(Zygote, Symbol("##794#795")))(::Nothing) at /home/henrique/.julia/packages/Zygote/bdE6T/src/lib/array.jl:37
 [3] (::getfield(Zygote, Symbol("##2078#back#796")){getfield(Zygote, Symbol("##794#795"))})(::Nothing) at /home/henrique/.julia/packages/ZygoteRules/Jn4LO/src/adjoint.jl:48
 [4] func_1 at ./In[16]:5 [inlined]
 [5] (::typeof(∂(func_1)))(::Array{Float64,1}) at /home/henrique/.julia/packages/Zygote/bdE6T/src/compiler/interface2.jl:0
 [6] #13 at ./In[17]:8 [inlined]
 [7] (::typeof(∂(λ)))(::Float64) at /home/henrique/.julia/packages/Zygote/bdE6T/src/compiler/interface2.jl:0
 [8] (::getfield(Zygote, Symbol("##32#33")){typeof(∂(λ))})(::Float64) at /home/henrique/.julia/packages/Zygote/bdE6T/src/compiler/interface.jl:38
 [9] gradient(::Function, ::Array{Float64,1}) at /home/henrique/.julia/packages/Zygote/bdE6T/src/compiler/interface.jl:47
 [10] jacobian_test(::typeof(func_1), ::Array{Float64,1}) at ./In[17]:8
 [11] top-level scope at In[19]:1

[sdewaele]

@ferrolho In general, you can use Zygote.Buffer to get rid of the mutating arrays error, see ?Zygote.Buffer. In this case, perhaps it is easier to write:

function func_1(x)
    x1 = 2x[1] + √x[3]
    x2 = 3x[1] + sin(x[4])
    x3 = 4x[1] + cos(x[5])
    dx = [x1,x2,x3]
    dx
end

Haven't tested gradient computation on this; I hope it works.

[ferrolho]

@sdewaele, thank you for replying. Indeed Zygote.Buffer is what I was looking for. Your solution also works but is not ideal for cases where the function is not limited to only a handful of variables.

For future reference, I leave the alternative version using Zygote.Buffer below:

function func_1(x)
    buf = Zygote.Buffer(x, 3)
    buf[1] = 2x[1] + √x[3]
    buf[2] = 3x[1] + sin(x[4])
    buf[3] = 4x[1] + cos(x[5])
    return copy(buf)
end

[jessebett]

What's the likelihood of @mcabbott's #414 solution being merged into the next release? Would love Zygote.jacobian to work out-of-the-box for the assignment I'm writing.

[cortner]

what is wrong with this:

Zygote.forward_jacobian(func_1, xx)[2]

??

[mkalia94]

@cortner For one, Zygote.forward_jacobian doesn't work when the function is a NN with elu or leakyrelu activation.

@mcabbott : have you tried Flux.train! with a jacobian piece in the loss? Very excited to have this in a future release!

[cossio]

Bump?