Jacobian differentiation with Enzyme

Specific Domains Numerics
,
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.