Adding A New AD

[willtebbutt]

Question❓

I'm interested in getting (the soon-to-be-renamed) Tapir.jl to work nicely with SciMLSensitivity.jl, as I understand this to be the way to get it to play nicely with the SciML ecosystem more broadly (please correct me if I'm wrong on this point!)

I can't see from the docs, or from a quick dive into some of the internals of this package, what the right way to go about this is. It looks like there's a few functions that I need to add methods for, but I'm not at all sure.

Could someone point me in the right direction?

[ChrisRackauckas]

You'd just define a _vecjacobian! overload: https://github.com/SciML/SciMLSensitivity.jl/blob/v7.66.2/src/derivative_wrappers.jl#L656-L752

along with the caches: https://github.com/SciML/SciMLSensitivity.jl/blob/v7.66.2/src/adjoint_common.jl#L212-L215

where that's used to build whatever caches are required to get a non-allocating vecjacobian.

You would probably want to opt it into callbacks https://github.com/SciML/SciMLSensitivity.jl/blob/v7.66.2/src/callback_tracking.jl because it supports mutation.

Similarly, GaussAdjoint needs a special cache

SciMLSensitivity.jl/src/gauss_adjoint.jl

Lines 416 to 432 in eded161

	elseif sensealg.autojacvec isa EnzymeVJP
	pf = let f = unwrappedf
	if DiffEqBase.isinplace(prob)
	function (out, u, _p, t)
	f(out, u, _p, t)
	nothing
	end
	else
	!DiffEqBase.isinplace(prob)
	function (out, u, _p, t)
	out .= f(u, _p, t)
	nothing
	end
	end
	end
	paramjac_config = zero(y), zero(y), Enzyme.make_zero(pf)
	pJ = nothing

and overload

SciMLSensitivity.jl/src/gauss_adjoint.jl

Lines 496 to 504 in eded161

	elseif sensealg.autojacvec isa EnzymeVJP
	tmp3, tmp4, tmp6 = paramjac_config
	vtmp4 = vec(tmp4)
	vtmp4 .= λ
	out .= 0
	Enzyme.autodiff(
	Enzyme.Reverse, Enzyme.Duplicated(pf, tmp6), Enzyme.Const,
	Enzyme.Duplicated(tmp3, tmp4),
	Enzyme.Const(y), Enzyme.Duplicated(p, out), Enzyme.Const(t))

and QuadratureAdjoint needs a fairly similar one

SciMLSensitivity.jl/src/quadrature_adjoint.jl

Lines 213 to 239 in eded161

	elseif sensealg.autojacvec isa EnzymeVJP
	pf = let f = unwrappedf
	if DiffEqBase.isinplace(prob) && prob isa RODEProblem
	function (out, u, _p, t, W)
	f(out, u, _p, t, W)
	nothing
	end
	elseif DiffEqBase.isinplace(prob)
	function (out, u, _p, t)
	f(out, u, _p, t)
	nothing
	end
	elseif !DiffEqBase.isinplace(prob) && prob isa RODEProblem
	function (out, u, _p, t, W)
	out .= f(u, _p, t, W)
	nothing
	end
	else
	!DiffEqBase.isinplace(prob)
	function (out, u, _p, t)
	out .= f(u, _p, t)
	nothing
	end
	end
	end
	paramjac_config = zero(y), zero(y), Enzyme.make_zero(pf)
	pJ = nothing

and

SciMLSensitivity.jl/src/quadrature_adjoint.jl

Lines 293 to 300 in eded161

	elseif sensealg.autojacvec isa EnzymeVJP
	tmp3, tmp4, tmp6 = paramjac_config
	tmp4 .= λ
	out .= 0
	Enzyme.autodiff(
	Enzyme.Reverse, Enzyme.Duplicated(pf, tmp6), Enzyme.Const,
	Enzyme.Duplicated(tmp3, tmp4),
	Enzyme.Const(y), Enzyme.Duplicated(p, out), Enzyme.Const(t))

[willtebbutt]

Lovely, thanks for this. Is there a recommended way to test that my overloads are implemented correctly?

[ChrisRackauckas]

You'll see that this file runs through lots of combinations:

https://github.com/SciML/SciMLSensitivity.jl/blob/master/test/adjoint.jl#L46-L196

If you add a vjp option for Tapir then you should be able to just add an autojacvec = AutoTapir() and add a test that it's matching. Then there's like 5 sets of them.

Because this grew organically over time it's a bit verbose, but we'll make that into a loop some day 😅

[willtebbutt]

Haha excellent -- thanks for the info!