Use DifferentiationInterface for AD in Implicit Solvers

[jClugstor]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

This is at a point where we can do stuff like this:

using OrdinaryDiffEqCore
using OrdinaryDiffEqSDIRK
using ADTypes
using EnzymeCore

function lorenz!(du, u, p, t)
    du[1] = 10.0 * (u[2] - u[1])
    du[2] = u[1] * (28.0 - u[3]) - u[2]
    du[3] = u[1] * u[2] - (8 / 3) * u[3]
end

u0 = [1.0; 0.0; 0.0]
tspan = (0.0, 100.0)
prob = ODEProblem{true, SciMLBase.NoSpecialize}(lorenz!, u0, tspan)
sol = solve(prob, ImplicitEuler(autodiff=AutoEnzyme(function_annotation = EnzymeCore.Const)))

and it actually uses sparsity detection and greedy jacobian coloring plus Enzyme to compute the Jacobians.

Some things I'm unsure about:

The current behavior is to use Jacobian coloring and SparseDiffTools by default. In order to keep that up, we have to wrap any ADType given in an AutoSparse unless it's already an AutoSparse. This does change the ADType that the user entered to be wrapped in an AutoSparse, which feels weird to me. Maybe there should be an option to just directly use the ADType entered, but by default we wrap it into an AutoSparse? I'm not sure.
The biggest issue is that the way the sparsity detectors work with DI is by using operator overloading (both TracerSparsityDetector and SymbolicsSparsityDetector do), but that's an issue when using AutoSpecialilzation, because of the FunctionWrappers. The solution I found was to just unwrap the function in the preparation process. I'm not sure what performance implication this will have, but I don't think it should do much, since the preparation should be run just once.
There's still pieces in here that use raw SparseDiffTools, (build_J_W) that I haven't looked in to how to convert to DI yet.
I may need to fix some of the versions.
There are some places that are getting sparse things where it's not expected.

[jClugstor]

In order for this to be completely done we'll need a DI equivalent for the SparseDiffTools JacVec operator that will be stored in the caches, for the W operators. I think I could just wrap the DI pushforward in an operator, but better to do a long term solution. In a recent Julia slack thread (https://julialang.slack.com/archives/C6G240ENA/p1735254065747829) there were a couple of solutions.

Is a good way to do this to make an extension in SciMLOperators for DifferentiationInterface that will have something like a DI_pushforward operator that basically wraps up the pushforward function in a FunctionOperator?

@ChrisRackauckas @oscardssmith @gdalle Any thoughts?

[ChrisRackauckas]

Is a good way to do this to make an extension in SciMLOperators for DifferentiationInterface that will have something like a DI_pushforward operator that basically wraps up the pushforward function in a FunctionOperator?

Yes

[ChrisRackauckas]

@avik-pal might already have one?

[gdalle]

Awesome work @jClugstor, thanks! Ping me when this is ready for a first round of DI-specific review.

This is at a point where we can do stuff like this [...] and it actually uses sparsity detection and greedy jacobian coloring plus Enzyme to compute the Jacobians.

Just to be clear, this wasn't possible before? So is this the first time that Enzyme can be used out-of-the-box to solve ODEs?

The current behavior is to use Jacobian coloring and SparseDiffTools by default. In order to keep that up, we have to wrap any ADType given in an AutoSparse unless it's already an AutoSparse. This does change the ADType that the user entered to be wrapped in an AutoSparse, which feels weird to me. Maybe there should be an option to just directly use the ADType entered, but by default we wrap it into an AutoSparse? I'm not sure.

Another option, which requires a bit more work (and is probably not worth it) would be to make SparseDiffTools compatible with the sparsity API of ADTypes v1. I think it might allow a more seamless upgrade. See e.g. JuliaDiff/SparseDiffTools.jl#298 for the detection aspect, and there should be a similar issue for the coloring aspect.

Speaking of SparseDiffTools, it still has an edge over DI when combined with FiniteDiff. The PR JuliaDiff/FiniteDiff.jl#191 could fix that, maybe @oscardssmith would be willing to take another look?

The biggest issue is that the way the sparsity detectors work with DI is by using operator overloading (both TracerSparsityDetector and SymbolicsSparsityDetector do), but that's an issue when using AutoSpecialization, because of the FunctionWrappers. The solution I found was to just unwrap the function in the preparation process. I'm not sure what performance implication this will have, but I don't think it should do much, since the preparation should be run just once.

Agreed, preparation is a one-time cost so I don't think we should worry too much (at least in the prototype stage).

There are some places that are getting sparse things where it's not expected.

What do you mean by unexpected sparse things? SparseMatrixCSC instead of Matrix? Can you give an example?

In order for this to be completely done we'll need a DI equivalent for the SparseDiffTools JacVec operator that will be stored in the caches, for the W operators. I think I could just wrap the DI pushforward in an operator, but better to do a long term solution. In a recent Julia slack thread (https://julialang.slack.com/archives/C6G240ENA/p1735254065747829) there were a couple of solutions.
Is a good way to do this to make an extension in SciMLOperators for DifferentiationInterface that will have something like a DI_pushforward operator that basically wraps up the pushforward function in a FunctionOperator?

We may also want to involve @oschulz and his AutoDiffOperators package to avoid duplication of efforts?

As a side note, DifferentiationInterface only has two dependencies: ADTypes and LinearAlgebra. For packages that use it extensively, I think it's reasonable to make it a full dep instead of a weakdep.

[gdalle]

Does Enzyme need to become a dependency? This adds significant install overhead, but if AutoEnzyme is to be the new default AD then it makes sense

[jClugstor]

Yeah, probably doesn't need to be a dependency unless we're committing to having it be the default.

[gdalle]

Suggested change

	DifferentiationInterface = "0.6.23"
	DifferentiationInterface = "0.6.28"

the other deps are also missing compat bounds?

[gdalle]

Suggested change

	DifferentiationInterface = "0.6.23"
	DifferentiationInterface = "0.6.31"

[gdalle]

Note that DI does not explicitly support complex numbers yet. What I mean by that is that we forward things to the backend as much as possible, so if the backend does support complex numbers then it will probably work, but there are no tests or hard API guarantees on that. See JuliaDiff/DifferentiationInterface.jl#646 for the discussion

[gdalle]

Also note that some differentiation operators are not defined unambiguously for complex numbers (e.g. the derivative for complex input)

[wsmoses]

Enzyme has an explicit variant of modes for complex numbers, that it probably would be wise to similarly wrap here (by default it will instead err warning about ambiguity if a function returns a complex number otherwise): https://enzyme.mit.edu/julia/stable/api/#EnzymeCore.ReverseHolomorphic . @gdalle I'm not sure DI supports this yet? so perhaps that means you may need to just call Enzyme.jacobian / autodiff directly in that case

[gdalle]

@jClugstor can you maybe specify where we will encounter complex numbers by filling the following table?

	derivative	jacobian
complex inputs possible	yes / no	yes / no
complex outputs possible	yes / no	yes / no

When there are both complex inputs and complex outputs, that's where we run into trouble because we cannot represent derivatives as a single scalar. In that case, the differentiation operators are not clearly defined (the Jacobian matrix is basically twice as big as it should be) so we would need to figure out what convention the ODE solvers need (see https://discourse.julialang.org/t/taking-complex-autodiff-seriously-in-chainrules/39317).

@wsmoses I understand your concern, but I find it encouraging that DI actually allowed Enzyme to be used here for the first time (or at least so I've been told). This makes me think that the right approach is to handle complex numbers properly in DI instead of introducing a special case for Enzyme?

[wsmoses]

sure adding proper complex number support to DI would be great, but a three line change here to use in-spec Complex support when there's already overloads for other ADTypes feels reasonable?

e.g. something like

function jacobian(f, x::AbstractArray{<:Complex}, integrator::WhatevertheTypeIs{<:AutoEnzyme})
  Enzyme.jacobian(ReverseHolomorphic, f, x)
end

from the discussion in JuliaDiff/DifferentiationInterface.jl#646 I think DI complex support is a much thornier issue. In particular, various tools have different conventions (e.g. jax vs pytorch pick different conjugates of what is propagated). So either DI needs to make a choice and shim/force all tools to use it (definitely doable), and then user code must be converted to that convention (e.g. a separate shim on the user side). For example, suppose DI picked a different conjugate from forwarddiff.jl. DI could write its shim once in forward diff to convert which is reasonable. But suppose one was defining a custom rule within ForwardDiff and the code called DI somewhere, now that user code needs to conditionally do a different the shim to conjugate which feels kind of nasty to be put everywhere (in contrast to a self consistent assumption). I suppose the other alternative is for DI to not pick a convention, but that again prevents users from using since it's not possible to know whether they get the correct value for them -- and worse, they won't know when they need to do a conversion or not.

Thus, if complex support is desired, a three line patch where things are explicitly supported seems okay (at least until the DI story is figured out)

[gdalle]

I agree that for now, this change seems to do the job (although it raises the question of consistency with the other backends that are handled via DI). But what will happen if the function in question is not holomorphic? That's the thorniest part of the problem, and that's why I wanted to inquire a bit more as to what kind of functions we can expect. Perhaps @jClugstor or @ChrisRackauckas can tell us more?

In any case, I have started a discussion on Discourse to figure out the right conventions: https://discourse.julialang.org/t/choosing-a-convention-for-complex-numbers-in-differentiationinterface/124433

[gdalle]

Also note that the Enzyme-specific fix only handles dense Jacobians, not sparse Jacobians (which are one of the main reasons to use DI in the first place)

[jClugstor]

Sorry, I can't really tell you much about the complex number support, other than previously only ForwardDiff or FiniteDiff were used, so when someone used an implicit solver on a complex problem, their conventions were used I guess. Also just wanted to note that the code this comment is on is just making sure that the FiniteDiff fdtype isn't complex if the function is a function wrapper and doesn't have to do with complex numbers through the solver in general.

[gdalle]

The latest release of DI inches closer to support for complex numbers. I read a little about conventions for non-holomorphic differentiation and it was a mess, so as a starting point DI assumes that the function is holomorphic. If you want e.g. a Jacobian, it is pretty much the only convention that makes sense anyway, otherwise you end up with a $2n \times 2n$ matrix. I have added complex holomorphic test scenarios, and right now FiniteDiff works on them. I'll test Enzyme soon enough.

[avik-pal]

@avik-pal might already have one?

Add a dispatch to https://github.com/SciML/NonlinearSolve.jl/blob/master/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl#L115

[jClugstor]

@gdalle

Just to be clear, this wasn't possible before? So is this the first time that Enzyme can be used out-of-the-box to solve ODEs?

As far as I know this is the first time Enzyme has been used for the implicit solvers yes.

[jClugstor]

@avik-pal I noticed that the constructors for your JacobianOperator take a NonlinearProblem , but doesn't use it that much. Would it make sense to create a constructor JacobianOperator(f::AbstractSciMLFunction, ...) that does essentially the same thing and put it in SciMLOperators?

[avik-pal]

@avik-pal I noticed that the constructors for your JacobianOperator take a NonlinearProblem , but doesn't use it that much. Would it make sense to create a constructor JacobianOperator(f::AbstractSciMLFunction, ...) that does essentially the same thing and put it in SciMLOperators?

the prepare_jvp and prepare_vjp functions assume a 2/3 arg function for oop/iip respectively, that won't hold for ordinarydiffeq

[gdalle]

Hey @jClugstor, just a friendly ping, do you need a hand on this one?

[jClugstor]

Not particularly, I've just been otherwise occupied the past few days.

I'm just going through all the tests and making sure that they pass one by one. Most of the difficulties have been related to making sure that the sparse AD works. For example making sure that if using sparse AD certain caches are constructed with sparse LU operators when needed etc.

Next step is making sure the stats tests pass, previously all of the function calls in the Jacobian/JacVec calculations were being counted. But obviously that's not true when using DI, so we'll need a better way to count how many times the user provided function has been called. I have some ideas already.

[gdalle]

I did something similar in DITest, you may want to take a look: https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/main/DifferentiationInterfaceTest/src/tests/benchmark.jl#L1-L22

[gdalle]

Also note that counting function executions only makes sense for AD backend like FiniteDiff and ForwardDiff. When you use Enzyme, it doesn't actually call the function while differentiating it (that's the difference between source transformation and operator overloading). So I'm not sure how you plan to do the counting in that case?

[ChrisRackauckas]

We count the number of times Enzyme is called.

[jClugstor]

I got pretty close to getting the stats working, but I was off by one for some of them. I think the reason is that when the DI prep object is created during the building of the jac_config, using the SparseMatrixColoring tools, it calls the function with tracers as the inputs. So in the tests that counts as a call to the function. But it's not possible to increment the nf stat when building the jacobian config, since we can't access the integrator at that point. So it's a little tricky.

Wrapping the user function so that there's a counter every time it's called would work for FiniteDiff, ForwardDiff, but it sounds like it won't work for Enzyme and Zygote etc.

Maybe one way is to have those counter wrappers have custom AD rules, so that when they're differentiated, it just differentiates the function inside, and iterates the counter?

[gdalle]

I think the reason is that when the DI prep object is created during the building of the jac_config, using the SparseMatrixColoring tools, it calls the function with tracers as the inputs.

Indeed that happens, but there might also be a call for reasons like getting the size and eltype of y = f(x). Do we really need to enforce a strict number of calls during preparation?

During execution too, the number of calls might be off by one, depending on the backend and how it handles things like value_and_jacobian

[jClugstor]

I think for now it's reasonable to disregard the function calls from the Jacobian preparation. Ideally though I think we would like to make sure we count every function call. I think that can be for a later time. @oscardssmith does that sound ok?

[oscardssmith]

sounds good.

[gdalle]

So what's next?

[jClugstor]

Once I get the stats tests working again, it looks like this will be really close. Many of the failing CI tests are also failing on main.

But:

• The Rosenbrock tests will need fixing
• Resizing for the jac_configs needs to be fixed, that might get tricky
• Many of the downstream integration tests will need to be fixed as well

[gdalle]

Resizing for the jac_configs needs to be fixed, that might get tricky

For this part specifically, I put a hidden preparation update function called prepare!_jacobian in DI, which can take an already existing prep object and work from it. At the moment this prep is simply discarded, but if you point me to the original code which does this config resizing, I will be able to port it to DI in order to save work with some backends.

[jClugstor]

The relevant functions are here:

OrdinaryDiffEq.jl/lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl

Line 326 in 013c00d

[gdalle]

What does it even mean to resize the Jacobian config for a sparse Jacobian? How do you anticipate the sparsity pattern for a bigger input variable 🤔 ?

[ChrisRackauckas]

Think about a PDE semi-discretization in 1D: it's just a banded matrix where the bands are all the same, it's just the number of times you repeat the pattern. This is very common for example in adaptive meshing PDE solvers.

[ChrisRackauckas]

Here's a stop gap solution:

1. Support resizing in dense cases
2. Support resizing in auto-sparse cases by triggering a new sparsity detection on resize! of the cache
3. Throw an informative error that for sparse matrices support will come in the future and point to an issue and see who shows up.

I think if auto-sparse is fast enough, (3) might go away naturally. I know the Trixi case only really uses the SSP methods so they would avoid this. Some of the finite element cases may be investigating (3) but we don't have full support already, so this basically all folds into the next OrdinaryDiffEq breaking v7 and we can just look forward not backwards on this. But I believe most finite element cases would appreciate (2) anyways?

The other thing that could be done is that the case of (3) would require that the new sparsity pattern is given with a special form of resize!. But we don't even have a public API way of doing that here (we have non-public ways...) so we can make the switch without such a thing existing.