JuliaDiff / DifferentiationInterface.jl

An interface to various automatic differentiation backends in Julia.
https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface
MIT License
202 stars 15 forks source link

Complex number support #646

Open ErikQQY opened 3 days ago

ErikQQY commented 3 days ago

Find this while implementing https://github.com/SciML/BoundaryValueDiffEq.jl/pull/258

This MWE is just proof of the idea and may not be meaningful, but I think it can still prove that the current DI lacks support for sparse Jacobian of complex numbers.

using ForwardDiff, DifferentiationInterface, SparseMatrixColorings, ADTypes
const DI = DifferentiationInterface
#backend = AutoForwardDiff() # this works
backend = AutoSparse(
    AutoForwardDiff(),
    sparsity_detector = ADTypes.KnownJacobianSparsityDetector(ones(3, 3)),
    coloring_algorithm = ConstantColoringAlgorithm(ones(3, 3), ones(Int64, 3))
)
u0 = [1.0; 0.0; 0.0] .+1im
jac_cache = DI.prepare_jacobian(nothing, similar(u0), backend, u0)

stack trace

ERROR: MethodError: no method matching DifferentiationInterfaceSparseMatrixColoringsExt.PushforwardSparseJacobianPrep(::DifferentiationInterface.BatchSizeSettings{…}, ::SparseMatrixColorings.ColumnColoringResult{…}, ::Matrix{…}, ::Vector{…}, ::Vector{…}, ::DifferentiationInterfaceForwardDiffExt.ForwardDiffTwoArgPushforwardPrep{…})

Closest candidates are:
  DifferentiationInterfaceSparseMatrixColoringsExt.PushforwardSparseJacobianPrep(::BS, ::C, ::M, ::S, ::R, ::E) where {BS<:DifferentiationInterface.BatchSizeSettings, C<:(AbstractColoringResult{:nonsymmetric, :column}), M<:(AbstractMatrix{<:Real}), S<:(AbstractVector{<:Tuple{Vararg{T, N}} where {N, T}}), R<:(AbstractVector{<:Tuple{Vararg{T, N}} where {N, T}}), E<:DifferentiationInterface.PushforwardPrep}
   @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/bulUW/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:11

Stacktrace:
 [1] _prepare_sparse_jacobian_aux_aux(::DifferentiationInterface.BatchSizeSettings{…}, ::SparseMatrixColorings.ColumnColoringResult{…}, ::Vector{…}, ::Tuple{…}, ::AutoSparse{…}, ::Vector{…})
   @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/bulUW/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:107
 [2] _prepare_sparse_jacobian_aux(::DifferentiationInterface.PushforwardFast, ::Vector{…}, ::Tuple{…}, ::AutoSparse{…}, ::Vector{…})
   @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/bulUW/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:81
 [3] prepare_jacobian(::Nothing, ::Vector{…}, ::AutoSparse{…}, ::Vector{…})
   @ DifferentiationInterfaceSparseMatrixColoringsExt ~/.julia/packages/DifferentiationInterface/bulUW/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:49
 [4] top-level scope
   @ ~/Random/test2.jl:27
Some type information was truncated. Use `show(err)` to see complete types.

I believe the culprit is the annotation of the compressed matrix being too restricted in the SparseMatrixColorings extensions, I locally changed this into more generic ones and the errors are gone.

https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/cc8818a2bb0fb3dab2abf29ba213f89213a8613a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl#L3-L33

gdalle commented 3 days ago

Note that DI currently doesn't guarantee anything vis a vis complex number support (this should probably be emphasized in the docs, PR welcome). The reason is that many backends don't support them at all, and the remaining ones probably have wildly diverging conventions. So while we can relax the type annotation, complex support is not part of the API at the moment (which also means it is not tested).

ErikQQY commented 3 days ago

I can understand that different AD backends may support this differently, so what is the current solution for this? I can PR to relax these type annotations and add proper tests(which I believe will resolve the ForwardDiff and ReverseDiff cases), we have to start somewhere right🤠?

gdalle commented 3 days ago

Feel free to open a PR replacing any <:Real annotation in the package with <:Number, since it will solve your immediate problem. But I'm not ready to add complex number support to the promises of DI's public API yet, precisely because we cannot guarantee the answers will be the same across backends. And the concepts themselves are not necessarily well defined. Do you have an idea which backends support complex numbers and how their support differs? I've never looked into it.

ErikQQY commented 3 days ago

From my testing, the funny thing is that the dense AD using DI would just work when dealing with complex numbers(only tested on ForwardDiff), it’s only sparse AD that suffers problems.

gdalle commented 3 days ago

Sure, but what I mean is that even if it "just works" (because DI forwards everything without constraining types), different backends may have different interpretations of what the complex derivatives mean. And I don't know if it is wise to expose these different interpretations under a common interface, possibly confusing users by letting them believe the operations are semantically the same. Does that make sense? There's a long thread which I never really read in detail about this, but I might have to dive back in: https://discourse.julialang.org/t/taking-complex-autodiff-seriously-in-chainrules/39317

ErikQQY commented 2 days ago

And I don't know if it is wise to expose these different interpretations under a common interface, possibly confusing users by letting them believe the operations are semantically the same. Does that make sense?

Yeah, that does make sense. So now I just relax the type annotations of the compressed matrix for my special use case then.