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 :confused: 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[7] = -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鈥檚 interface.

Abstract differentiation didn鈥檛 add Enzyme yet.

In any case high level jacobian wrappers via Enzyme鈥檚 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[2], v[1]*v[1], v[1]*v[1]*v[1]]
    end

	jac = Enzyme.revjacobian(inout, [2.0, 3.0], Val(1); n_outs=Val(3))	
	@test length(jac) == 3
	@test jac[1] 鈮 [ 0.0, 1.0]
	@test jac[2] 鈮 [ 4.0, 0.0]
	@test jac[3] 鈮 [12.0, 0.0]
	
	jac = Enzyme.fwdjacobian(inout, [2.0, 3.0], Val(1))
	@test length(jac) == 2
	@test jac[1] 鈮 [ 0.0,  4.0, 12.0]
	@test jac[2] 鈮 [ 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[1] 鈮 [ 0.0, 1.0]
	@test jac[2] 鈮 [ 4.0, 0.0]
	@test jac[3] 鈮 [12.0, 0.0]
	
	jac = Enzyme.fwdjacobian(inout, [2.0, 3.0], Val(2))
	@test length(jac) == 2
	@test jac[1] 鈮 [ 0.0,  4.0, 12.0]
	@test jac[2] 鈮 [ 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 :slight_smile: