Jacobian differentiation with Enzyme

Specific Domains Numerics
,
1

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?

2

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

3

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 (https://github.com/EnzymeAD/Enzyme.jl/pull/280) 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):

https://github.com/EnzymeAD/Enzyme.jl/blob/97da73ddaa869bd12d4b538f2bec4d3a2e020df2/test/runtests.jl#L902

@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 (https://github.com/JuliaGPU/GPUCompiler.jl/pull/320) which hopefully will be merged shortly.

4

Awesome, thanks! I am eagerly awaiting the merge :slight_smile: