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]

Within DI, the function prepare!_jacobian(f!, y, prep, backend, x) allows an existing preparation output prep to be resized and adapted to the new input x. By default, it re-runs the preparation from scratch (which, in the sparse case, includes sparsity detection). But we can override it at will if we have faster ways to resize. I have only done it for ForwardDiff yet, but FiniteDiff will soon follow (as soon as FiniteDiff.jl#191 is done). Here's how it looks:

https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/5dfd7adec430c71a63f527a61962d5e5567e6702/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl#L372-L389

Can you take it out for a spin, and tell me whether anything else is missing from DI in your opinion?

[jClugstor]

I just finished up fixing the resizing, looks like it's working.

Up next, in-place DAE's aren't working with the tracing sparsity detection, due to the way the DAEResidualJacobianWrapper is set up.

[gdalle]

That's probably an issue for SparseConnectivityTracer to solve. But note that in the meantime, you can still use other methods for sparsity detection, like the good old Symbolics. Just switch the sparsity_detector argument in AutoSparse to Symbolics.SymbolicsSparsityDetector() and voilà!

[jClugstor]

Any detector based on operator overloading will fail I believe, it's just that the DAEResidualJacobianWrapper caches are built with Float64, and evaluating the function wrapped in them uses the caches.

[gdalle]

What did you use before?

• If you know the pattern in advance, there is ADTypes.KnownJacobianSparsityDetector
• If the problem is reasonably-sized, there's always DI.DenseSparsityDetector

[jClugstor]

Before I believe you had to supply a sparsity pattern to the ODEFunction manually, and then this function would find the colors or use the supplied colorvec, and pass it to the ForwardColorJacCache. So there wasn't any automatic sparsity detection.

[gdalle]

Okay then it's easy to at least keep the current behavior using the KnownJacobianSparsityDetector

[jClugstor]

Yes, I have it so if the user supplies a sparsity pattern KnownJacobianSparsityDetector is used, plus if a colorvec is supplied it uses that in a ConstantColoringAlgorithm.

[gdalle]

Seems like you got this covered

[gdalle]

As of the latest version of DI, FiniteDiff's sparse Jacobians use a new JVPCache (courtesy of @oscardssmith) and should therefore be non-allocating

[jClugstor]

I think I'm going to split this PR in to two parts, now that I understand what the previous behavior was better. Most of the existing problems stem from having the automatic sparsity detection apply automatically, where as actually previous behavior only used any kind of matrix coloring when a sparsity pattern was supplied to the ODEFunction.

So I think I'll split them up like this:

1. Make sure the previous behavior is preserved while using DI for all of the AD. This is basically done, every AD operation in there uses DI at this point, and we can look to see if a sparsity pattern and colorvec is supplied then use the DI mechanisms to take advantage of them. Also, because I already made a couple of changes to get the sparsity detectors to work through it all, if the user supplies an ADType AutoSparse with automatic detection, it should work in many cases already.
2. Allow for automatically applying sparsity and sparsity detection to the ADTypes, when asked for. This might require an addition to the API, maybe just a keyword argument? I'm thinking it would go something like this: if the user tells it to try automatic sparsity detection, and supplied a dense ADType, we try to build the Jacobian configs with automatic sparsity detection with TracerSparsityDetection, if that errors, we can try SymbolicSparsityDetection, and so on. This will need some rearranging of some stuff to make sure that the alg with the correct ADType gets supplied to the integrator.

[jClugstor]

@gdalle we'll need support for the (undocumented) dir keyword argument for FiniteDiff.jl. I think we'll just need to add it to the FiniteDiff ADType, and then we should be able to just plug it in to finite_difference_derivative etc. in DI.

[gdalle]

This might be a good opportunity to document it in FiniteDiff as well? Along with other args or kwargs we want to use?