# Jacobian differentiation with Enzyme

Hi all,

I am having some trouble understanding the docs for Enzyme.jl. Specifically, how to differentiate functions that return arrays and take arrays as inputs (i.e. how to get the Jacobian). I have tried a number of approaches, but they tend to crash Julia Could someone please show me how to get the jacobian of `sparse_F` as shown below in the code?

``````using Enzyme, SparseArrays, LinearAlgebra

# setup data
θ = rand(3)
x = rand(11)
ν = rand(10)
λ = rand(24)
S = sprand(Float64, 6, 7, 0.1)
Cp = sprand(Float64, 1, 11, 0.4)
_h = rand(24)
_d = spzeros(10)
_c = spzeros(11)
λ = zeros(24)
_Q = spzeros(11, 11)
_c = -1.0
θ = [30.0, 70.0, 90.0]
x = [1.9873817034700316, 3.9747634069400632, 1.9873817034700316, 1.9873817034700316, 1.9873817034700316, 3.9747634069400632, 1.9873817034700316, 0.028391167192429023, 0.022082018927444796, 0.13249211356466878, 0.13249211356466878]
ν = [ 0.0, 0.0, 0.028391167192429023, 0.24290220820189273,0.7287066246056783, 0.0, 1.9873817034700316, 3.9747634069400632, 5.962145110410095, 7.9495268138801265,]
λ[end] = -1.9873817034700316
S = [
1.0  0.0  0.0  0.0  0.0  0.0  -1.0
0.0  1.0  0.0  0.0  0.0  0.0  -2.0
0.0  0.0  1.0  0.0  0.0  0.0  -1.0
0.0  0.0  0.0  1.0  0.0  0.0  -1.0
0.0  0.0  0.0  0.0  1.0  0.0  -1.0
0.0  0.0  0.0  0.0  0.0  1.0  -2.0
]
se_vars = ((2, -1.0), (3, -1.0),(3, -3.0), (1, -2.0), (1, -1.0),)

# create functions
Se(θ) = sparse(
[1, 2, 3, 4, 3],
[3, 4, 4, 5, 5],
[k / θ[idx] for (idx, k) in se_vars],
4,
7,
)

E(θ) = [
S spzeros(6, 4)
Se(θ) I(4)
]

d = _ -> _d

c = _ -> _c

M = _ -> [
-I(11)
I(11)
-Cp
Cp
]

h = _ -> _h

Q = _ -> _Q

xidxs = 1:length(x)
νidxs = last(xidxs) .+ (1:length(ν))
λidxs = last(νidxs) .+ (1:length(λ))
θidxs = last(λidxs) .+ (1:length(θ))

sparse_F(z) = [
Q(z[θidxs]) * z[xidxs] + c(z[θidxs]) -
E(z[θidxs])' * z[νidxs] - M(z[θidxs])' * z[λidxs]
E(z[θidxs]) * z[xidxs] - d(z[θidxs])
spdiagm(z[λidxs]) * (M(z[θidxs]) * z[xidxs] - h(z[θidxs]))
]

z = [x; ν; λ; θ]
sparse_F(z) # this function should be differentiated wrt z
``````

I have tried many variations of the following, but it always crashes Julia:

``````function f(out, z)
out .= sparse_F(z)
nothing # the docs say this you should do this for array output funcs
end

out = spzeros(45)
d_out = spzeros(45)
d_z = zeros(length(z))
autodiff(f, Active, Duplicated(out, d_out), Duplicated(z, d_z))
``````

Does anybody have an idea?

It should just work via AbstractDifferentiation.jl’s interface.

Abstract differentiation didn’t add Enzyme yet.

In any case high level jacobian wrappers via Enzyme’s forward and reverse mode jacobians are coming in this PR (Handle vector width abi by wsmoses · Pull Request #280 · EnzymeAD/Enzyme.jl · GitHub) and eventually can be called as follows (the `Val(1)` is a chunk size (anywhere `1` to length(out) is fine and when doing reverse mode `n_outs` is the length of the output):

``````@testset "Jacobian" begin
function inout(v)
[v, v*v, v*v*v]
end

jac = Enzyme.revjacobian(inout, [2.0, 3.0], Val(1); n_outs=Val(3))
@test length(jac) == 3
@test jac ≈ [ 0.0, 1.0]
@test jac ≈ [ 4.0, 0.0]
@test jac ≈ [12.0, 0.0]

jac = Enzyme.fwdjacobian(inout, [2.0, 3.0], Val(1))
@test length(jac) == 2
@test jac ≈ [ 0.0,  4.0, 12.0]
@test jac ≈ [ 1.0,  0.0,  0.0]

jac = Enzyme.revjacobian(inout, [2.0, 3.0], Val(2); n_outs=Val(3))
@test length(jac) == 3
@test jac ≈ [ 0.0, 1.0]
@test jac ≈ [ 4.0, 0.0]
@test jac ≈ [12.0, 0.0]

jac = Enzyme.fwdjacobian(inout, [2.0, 3.0], Val(2))
@test length(jac) == 2
@test jac ≈ [ 0.0,  4.0, 12.0]
@test jac ≈ [ 1.0,  0.0,  0.0]
end
``````

Note that that PR requires (Use function type in fspec instead of value by wsmoses · Pull Request #320 · JuliaGPU/GPUCompiler.jl · GitHub) which hopefully will be merged shortly.

Awesome, thanks! I am eagerly awaiting the merge 