What would we need to become backend independent via using DifferentiationInterface?

[oxinabox]

This is a space for a discussion on what would be required to change here and in DifferentiationInterface to make this package backend agnostic.
(it equally could have been made on DifferentiationInterface.jl but I had to pick a location)

I had thought we were a fair way away from that, but i see DifferentiationInterface is already supporting caches and is using ADTypes.

Very likely this would still be a breaking change for SciMLSensitivity, but hopefully not a major one.
I think it would probably be worth it.
@ChrisRackauckas once referred to SciMLSensitivity as "the final boss of julia", because it depends on so many packages. Including ones like Enzyme and Zygote which are sensitive to changes to julias internals, and which sometimes break during precompilation because of this, which makes SciMLSensitivity break on nightly even if you are not using those packages.

cc @gdalle

[gdalle]

Happy to discuss anything that might be needed! Note that it would probably help reduce latency on the user side to avoid loading every backend at once.

Cc @adrhill

[ChrisRackauckas]

Each vjp just has a dispatch. For example, the Zygote one is probably the most basic:

https://github.com/SciML/SciMLSensitivity.jl/blob/master/src/derivative_wrappers.jl#L514-L647

I would say, start by moving them one by one to DifferentiationInterface.jl. I wouldn't try to do them all in one go, Enzyme in particular you will have issues making DifferentiationInterface.jl be performant in its current form, so just try and make a DI dispatch, and remove dispatches one at a time as they are benchmarked to be performant. I think the first few would take only a few hours.

[gdalle]

What do you think is needed to make Enzyme perform well with DifferentiationInterface? Is it just adding more overloads to the extension, or something deeper?

[ChrisRackauckas]

It has some specialized cache handling https://github.com/SciML/SciMLSensitivity.jl/blob/master/src/derivative_wrappers.jl#L658-L681 and differentiates multiple arguments simultaneously (https://github.com/SciML/SciMLSensitivity.jl/blob/master/src/derivative_wrappers.jl#L696-L723). I don't see an interface for that in DI

[gdalle]

I'll take a look and make some suggestions

[gdalle]

Coming back to this old issue, DI now supports multiple arguments (albeit not multiple active arguments). Would that be sufficient to get started?

[ChrisRackauckas]

That would be enough to get the WIP PR up, but it won't be mergable until we can differentiate both u and p at the same time.