# Jacobian differentiation with Enzyme

**URL:** <https://discourse.julialang.org/t/jacobian-differentiation-with-enzyme/79912>\
**Category:** Numerics\
**Tags:** differentiation, enzyme\
**Created:** [April 23, 2022, 8:22pm UTC](https://discourse.julialang.org/t/jacobian-differentiation-with-enzyme/79912 "2022-04-23T20:22:35Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Elmo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elmo/32/17979_2.png) [@Elmo](https://discourse.julialang.org/u/Elmo)\
**Post date:** [April 23, 2022, 8:22pm UTC](https://discourse.julialang.org/t/jacobian-differentiation-with-enzyme/79912/1 "2022-04-23T20:22:35Z")

</div>

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?

```julia
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:

```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?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [April 23, 2022, 10:32pm UTC](https://discourse.julialang.org/t/jacobian-differentiation-with-enzyme/79912/2 "2022-04-23T22:32:29Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [April 24, 2022, 2:48am UTC](https://discourse.julialang.org/t/jacobian-differentiation-with-enzyme/79912/3 "2022-04-24T02:48:25Z")

</div>

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](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](https://github.com/EnzymeAD/Enzyme.jl/blob/97da73ddaa869bd12d4b538f2bec4d3a2e020df2/test/runtests.jl#L902)

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

---

<div class="post-metadata">

**Author:** ![Elmo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/elmo/32/17979_2.png) [@Elmo](https://discourse.julialang.org/u/Elmo)\
**Post date:** [April 24, 2022, 8:57am UTC](https://discourse.julialang.org/t/jacobian-differentiation-with-enzyme/79912/4 "2022-04-24T08:57:28Z")

</div>

Awesome, thanks! I am eagerly awaiting the merge 🙂