feat: allow and test holomorphic derivatives (#687)

gdalle · web-flow · commit 66f7ff3198b8 · 2025-01-21T14:39:28.000+01:00
* feat: allow and test holomorphic derivatives

* API doc

* Add tests

* Skip matrix tests for FastDifferentiation

* Fix

* Proper use of dot

* Dep
diff --git a/.gitignore b/.gitignore
@@ -10,4 +10,5 @@
 
 *.csv
 
-playground.jl
+playground.jl
+.vscode
diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml
@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.30"
+version = "0.6.31"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/DifferentiationInterface/docs/src/faq/limitations.md b/DifferentiationInterface/docs/src/faq/limitations.md
@@ -7,5 +7,7 @@ As a result, DifferentiationInterface only supports a single active argument, ca
 
 ## Complex numbers
 
-At the moment, complex numbers are only handled by a few AD backends, sometimes using different conventions.
-As a result, DifferentiationInterface is only tested on real numbers and complex number support is not part of its API guarantees.
+Complex derivatives are only handled by a few AD backends, sometimes using different conventions.
+To find the easiest common ground, DifferentiationInterface assumes that whenever complex numbers are involved, the function to differentiate is holomorphic.
+This functionality is still considered experimental and not yet part of the public API guarantees.
+If you work with non-holomorphic functions, you will need to manually separate real and imaginary parts.
diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl
@@ -319,6 +319,9 @@ function _sparse_jacobian_aux!(
         )
 
         for b in eachindex(batched_results[a])
+            if eltype(x) <: Complex
+                batched_results[a][b] .= conj.(batched_results[a][b])
+            end
             copyto!(
                 view(compressed_matrix, 1 + ((a - 1) * B + (b - 1)) % N, :),
                 vec(batched_results[a][b]),
diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl b/DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl
@@ -220,6 +220,9 @@ function _sparse_jacobian_aux!(
         )
 
         for b in eachindex(batched_results_reverse[a])
+            if eltype(x) <: Complex
+                batched_results_reverse[a][b] .= conj.(batched_results_reverse[a][b])
+            end
             copyto!(
                 view(compressed_matrix_reverse, 1 + ((a - 1) * Br + (b - 1)) % Nr, :),
                 vec(batched_results_reverse[a][b]),
diff --git a/DifferentiationInterface/src/first_order/derivative.jl b/DifferentiationInterface/src/first_order/derivative.jl
@@ -74,14 +74,14 @@ end
 function prepare_derivative(
     f::F, backend::AbstractADType, x, contexts::Vararg{Context,C}
 ) where {F,C}
-    pushforward_prep = prepare_pushforward(f, backend, x, (realone(x),), contexts...)
+    pushforward_prep = prepare_pushforward(f, backend, x, (one(x),), contexts...)
     return PushforwardDerivativePrep(pushforward_prep)
 end
 
 function prepare_derivative(
     f!::F, y, backend::AbstractADType, x, contexts::Vararg{Context,C}
 ) where {F,C}
-    pushforward_prep = prepare_pushforward(f!, y, backend, x, (realone(x),), contexts...)
+    pushforward_prep = prepare_pushforward(f!, y, backend, x, (one(x),), contexts...)
     return PushforwardDerivativePrep(pushforward_prep)
 end
 
@@ -95,7 +95,7 @@ function value_and_derivative(
     contexts::Vararg{Context,C},
 ) where {F,C}
     y, ty = value_and_pushforward(
-        f, prep.pushforward_prep, backend, x, (realone(x),), contexts...
+        f, prep.pushforward_prep, backend, x, (one(x),), contexts...
     )
     return y, only(ty)
 end
@@ -109,7 +109,7 @@ function value_and_derivative!(
     contexts::Vararg{Context,C},
 ) where {F,C}
     y, _ = value_and_pushforward!(
-        f, (der,), prep.pushforward_prep, backend, x, (realone(x),), contexts...
+        f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...
     )
     return y, der
 end
@@ -121,7 +121,7 @@ function derivative(
     x,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    ty = pushforward(f, prep.pushforward_prep, backend, x, (realone(x),), contexts...)
+    ty = pushforward(f, prep.pushforward_prep, backend, x, (one(x),), contexts...)
     return only(ty)
 end
 
@@ -133,7 +133,7 @@ function derivative!(
     x,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    pushforward!(f, (der,), prep.pushforward_prep, backend, x, (realone(x),), contexts...)
+    pushforward!(f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...)
     return der
 end
 
@@ -148,7 +148,7 @@ function value_and_derivative(
     contexts::Vararg{Context,C},
 ) where {F,C}
     y, ty = value_and_pushforward(
-        f!, y, prep.pushforward_prep, backend, x, (realone(x),), contexts...
+        f!, y, prep.pushforward_prep, backend, x, (one(x),), contexts...
     )
     return y, only(ty)
 end
@@ -163,7 +163,7 @@ function value_and_derivative!(
     contexts::Vararg{Context,C},
 ) where {F,C}
     y, _ = value_and_pushforward!(
-        f!, y, (der,), prep.pushforward_prep, backend, x, (realone(x),), contexts...
+        f!, y, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...
     )
     return y, der
 end
@@ -176,7 +176,7 @@ function derivative(
     x,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    ty = pushforward(f!, y, prep.pushforward_prep, backend, x, (realone(x),), contexts...)
+    ty = pushforward(f!, y, prep.pushforward_prep, backend, x, (one(x),), contexts...)
     return only(ty)
 end
 
@@ -189,9 +189,7 @@ function derivative!(
     x,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    pushforward!(
-        f!, y, (der,), prep.pushforward_prep, backend, x, (realone(x),), contexts...
-    )
+    pushforward!(f!, y, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...)
     return der
 end
 
diff --git a/DifferentiationInterface/src/first_order/jacobian.jl b/DifferentiationInterface/src/first_order/jacobian.jl
@@ -319,6 +319,9 @@ function _jacobian_aux(
     dx_batch = pullback(
         f_or_f!y..., pullback_prep, backend, x, only(batched_seeds), contexts...
     )
+    if eltype(x) <: Complex
+        dx_batch = map(conj, dx_batch)
+    end
     block = stack_vec_row(dx_batch)
     if aligned
         return block
@@ -345,6 +348,9 @@ function _jacobian_aux(
         dx_batch = pullback(
             f_or_f!y..., pullback_prep_same, backend, x, batched_seeds[a], contexts...
         )
+        if eltype(x) <: Complex
+            dx_batch = map(conj, dx_batch)
+        end
         block = stack_vec_row(dx_batch)
         if !aligned && a == A
             return block[1:B_last, :]
@@ -403,7 +409,7 @@ function _jacobian_aux!(
     (; N) = batch_size_settings
 
     pullback_prep_same = prepare_pullback_same_point(
-        f_or_f!y..., prep.pullback_prep, backend, x, batched_seeds[1], contexts...
+        f_or_f!y..., pullback_prep, backend, x, batched_seeds[1], contexts...
     )
 
     for a in eachindex(batched_seeds, batched_results)
@@ -418,6 +424,9 @@ function _jacobian_aux!(
         )
 
         for b in eachindex(batched_results[a])
+            if eltype(x) <: Complex
+                batched_results[a][b] .= conj.(batched_results[a][b])
+            end
             copyto!(
                 view(jac, 1 + ((a - 1) * B + (b - 1)) % N, :), vec(batched_results[a][b])
             )
diff --git a/DifferentiationInterface/src/first_order/pullback.jl b/DifferentiationInterface/src/first_order/pullback.jl
@@ -156,7 +156,7 @@ function _pullback_via_pushforward(
     contexts::Vararg{Context,C},
 ) where {F,C}
     t1 = pushforward(f, pushforward_prep, backend, x, (one(x),), contexts...)
-    dx = dot(dy, only(t1))
+    dx = dot(only(t1), dy)
     return dx
 end
 
@@ -170,7 +170,7 @@ function _pullback_via_pushforward(
 ) where {F,C}
     dx = map(CartesianIndices(x)) do j
         t1 = pushforward(f, pushforward_prep, backend, x, (basis(backend, x, j),), contexts...)
-        dot(dy, only(t1))
+        dot(only(t1), dy)
     end
     return dx
 end
@@ -241,7 +241,7 @@ function _pullback_via_pushforward(
     contexts::Vararg{Context,C},
 ) where {F,C}
     t1 = pushforward(f!, y, pushforward_prep, backend, x, (one(x),), contexts...)
-    dx = dot(dy, only(t1))
+    dx = dot(only(t1), dy)
     return dx
 end
 
@@ -258,7 +258,7 @@ function _pullback_via_pushforward(
         t1 = pushforward(
             f!, y, pushforward_prep, backend, x, (basis(backend, x, j),), contexts...
         )
-        dot(dy, only(t1))
+        dot(only(t1), dy)
     end
     return dx
 end
diff --git a/DifferentiationInterface/src/first_order/pushforward.jl b/DifferentiationInterface/src/first_order/pushforward.jl
@@ -158,7 +158,7 @@ function _pushforward_via_pullback(
     contexts::Vararg{Context,C},
 ) where {F,C}
     t1 = pullback(f, pullback_prep, backend, x, (one(y),), contexts...)
-    dy = dot(dx, only(t1))
+    dy = dot(only(t1), dx)
     return dy
 end
 
@@ -173,7 +173,7 @@ function _pushforward_via_pullback(
 ) where {F,C}
     dy = map(CartesianIndices(y)) do i
         t1 = pullback(f, pullback_prep, backend, x, (basis(backend, y, i),), contexts...)
-        dot(dx, only(t1))
+        dot(only(t1), dx)
     end
     return dy
 end
@@ -245,7 +245,7 @@ function _pushforward_via_pullback(
 ) where {F,C}
     dy = map(CartesianIndices(y)) do i  # preserve shape
         t1 = pullback(f!, y, pullback_prep, backend, x, (basis(backend, y, i),), contexts...)
-        dot(dx, only(t1))
+        dot(only(t1), dx)
     end
     return dy
 end
diff --git a/DifferentiationInterface/src/misc/from_primitive.jl b/DifferentiationInterface/src/misc/from_primitive.jl
@@ -1,17 +1,17 @@
 abstract type FromPrimitive <: AbstractADType end
 
-function basis(fromprim::FromPrimitive, x::AbstractArray{<:Real}, i)
+function basis(fromprim::FromPrimitive, x::AbstractArray, i)
     return basis(fromprim.backend, x, i)
 end
 
-function multibasis(fromprim::FromPrimitive, x::AbstractArray{<:Real}, inds)
+function multibasis(fromprim::FromPrimitive, x::AbstractArray, inds)
     return multibasis(fromprim.backend, x, inds)
 end
 
 check_available(fromprim::FromPrimitive) = check_available(fromprim.backend)
 inplace_support(fromprim::FromPrimitive) = inplace_support(fromprim.backend)
 
-function BatchSizeSettings(fromprim::FromPrimitive, x::AbstractArray{<:Real})
+function BatchSizeSettings(fromprim::FromPrimitive, x::AbstractArray)
     return BatchSizeSettings(fromprim.backend, x)
 end
 
diff --git a/DifferentiationInterface/src/utils/basis.jl b/DifferentiationInterface/src/utils/basis.jl
@@ -46,7 +46,7 @@ Construct the `i`-th standard basis array in the vector space of `a` with elemen
 If an AD backend benefits from a more specialized basis array implementation,
 this function can be extended on the backend type.
 """
-basis(::AbstractADType, a::AbstractArray{<:Real}, i) = basis(a, i)
+basis(::AbstractADType, a::AbstractArray, i) = basis(a, i)
 
 """
     multibasis(backend, a::AbstractArray, inds::AbstractVector)
@@ -58,18 +58,16 @@ Construct the sum of the `i`-th standard basis arrays in the vector space of `a`
 If an AD backend benefits from a more specialized basis array implementation,
 this function can be extended on the backend type.
 """
-multibasis(::AbstractADType, a::AbstractArray{<:Real}, inds) = multibasis(a, inds)
+multibasis(::AbstractADType, a::AbstractArray, inds) = multibasis(a, inds)
 
-function basis(a::AbstractArray{T,N}, i) where {T<:Real,N}
+function basis(a::AbstractArray{T,N}, i) where {T,N}
     return zero(a) + OneElement(i, one(T), a)
 end
 
-function multibasis(a::AbstractArray{T,N}, inds::AbstractVector) where {T<:Real,N}
+function multibasis(a::AbstractArray{T,N}, inds::AbstractVector) where {T,N}
     seed = zero(a)
     for i in inds
         seed += OneElement(i, one(T), a)
     end
     return seed
 end
-
-realone(x::Real) = one(x)
diff --git a/DifferentiationInterface/test/Back/FiniteDiff/test.jl b/DifferentiationInterface/test/Back/FiniteDiff/test.jl
@@ -2,6 +2,7 @@ using Pkg
 Pkg.add("FiniteDiff")
 
 using DifferentiationInterface, DifferentiationInterfaceTest
+using DifferentiationInterface: DenseSparsityDetector
 using FiniteDiff: FiniteDiff
 using SparseMatrixColorings
 using Test
@@ -23,10 +24,15 @@ test_differentiation(
     logging=LOGGING,
 );
 
-@testset verbose = true "Complex number support" begin
-    backend = AutoSparse(AutoFiniteDiff(); coloring_algorithm=GreedyColoringAlgorithm())
-    x = float.(1:3) .+ im
-    @test_skip @test_nowarn jacobian(identity, backend, x)
-    @test_skip @test_nowarn jacobian(copyto!, similar(x), backend, x)
-    @test_skip @test_nowarn hessian(sum, backend, x)
-end
+@testset "Complex" begin
+    test_differentiation(AutoFiniteDiff(), complex_scenarios(); logging=LOGGING)
+    test_differentiation(
+        AutoSparse(
+            AutoFiniteDiff();
+            sparsity_detector=DenseSparsityDetector(AutoFiniteDiff(); atol=1e-5),
+            coloring_algorithm=GreedyColoringAlgorithm(),
+        ),
+        complex_sparse_scenarios();
+        logging=LOGGING,
+    )
+end;
diff --git a/DifferentiationInterface/test/Back/Zygote/test.jl b/DifferentiationInterface/test/Back/Zygote/test.jl
@@ -42,6 +42,10 @@ end
     )
 end
 
+test_differentiation(
+    AutoZygote(), complex_scenarios(); excluded=[:gradient, :jacobian], logging=LOGGING
+);
+
 ## Sparse
 
 @testset "Sparse" begin
diff --git a/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl b/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl
@@ -1,6 +1,7 @@
 using DifferentiationInterface, DifferentiationInterfaceTest
-using DifferentiationInterface: AutoSimpleFiniteDiff, AutoReverseFromPrimitive
-using DifferentiationInterfaceTest
+using DifferentiationInterface:
+    AutoSimpleFiniteDiff, AutoReverseFromPrimitive, DenseSparsityDetector
+using SparseMatrixColorings
 using Test
 
 LOGGING = get(ENV, "CI", "false") == "false"
@@ -41,6 +42,8 @@ end
         default_scenarios(; include_constantified=true);
         logging=LOGGING,
     )
+
+    test_differentiation(backends, complex_scenarios(); logging=LOGGING)
 end
 
 @testset "Sparse" begin
@@ -59,6 +62,20 @@ end
         logging=LOGGING,
     )
 
+    @testset "Complex numbers" begin
+        test_differentiation(
+            AutoSparse.(
+                vcat(
+                    adaptive_backends, MixedMode(adaptive_backends[1], adaptive_backends[2])
+                );
+                sparsity_detector=DenseSparsityDetector(AutoSimpleFiniteDiff(); atol=1e-5),
+                coloring_algorithm=GreedyColoringAlgorithm(),
+            ),
+            complex_sparse_scenarios();
+            logging=LOGGING,
+        )
+    end
+
     @testset "SparseMatrixColorings access" begin
         jac_for_prep = prepare_jacobian(copy, MyAutoSparse(adaptive_backends[1]), rand(10))
         jac_rev_prep = prepare_jacobian(copy, MyAutoSparse(adaptive_backends[2]), rand(10))
diff --git a/DifferentiationInterfaceTest/Project.toml b/DifferentiationInterfaceTest/Project.toml
@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.9.1"
+version = "0.9.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -93,4 +93,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "DifferentiationInterface", "ExplicitImports", "FiniteDifferences", "Flux", "ForwardDiff", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StaticArrays", "Test", "Zygote"]
+test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "DifferentiationInterface", "ExplicitImports", "FiniteDiff", "FiniteDifferences", "Flux", "ForwardDiff", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StaticArrays", "Test", "Zygote"]
diff --git a/DifferentiationInterfaceTest/docs/src/api.md b/DifferentiationInterfaceTest/docs/src/api.md
@@ -35,6 +35,8 @@ sparse_scenarios
 component_scenarios
 gpu_scenarios
 static_scenarios
+complex_scenarios
+complex_sparse_scenarios
 ```
 
 ## Internals
diff --git a/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl b/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl
diff --git a/DifferentiationInterfaceTest/src/scenarios/complex.jl b/DifferentiationInterfaceTest/src/scenarios/complex.jl
diff --git a/DifferentiationInterfaceTest/src/scenarios/default.jl b/DifferentiationInterfaceTest/src/scenarios/default.jl
diff --git a/DifferentiationInterfaceTest/test/standard.jl b/DifferentiationInterfaceTest/test/standard.jl

Original file line number	Diff line number	Diff line change
`@@ -319,6 +319,9 @@ function _sparse_jacobian_aux!(`
`319`	`319`	`)`
`320`	`320`
`321`	`321`	`for b in eachindex(batched_results[a])`
	`322`	`+ if eltype(x) <: Complex`
	`323`	`+ batched_results[a][b] .= conj.(batched_results[a][b])`
	`324`	`+ end`
`322`	`325`	`copyto!(`
`323`	`326`	`view(compressed_matrix, 1 + ((a - 1) * B + (b - 1)) % N, :),`
`324`	`327`	`vec(batched_results[a][b]),`
Original file line number	Diff line number	Diff line change
`@@ -220,6 +220,9 @@ function _sparse_jacobian_aux!(`
`220`	`220`	`)`
`221`	`221`
`222`	`222`	`for b in eachindex(batched_results_reverse[a])`
	`223`	`+ if eltype(x) <: Complex`
	`224`	`+ batched_results_reverse[a][b] .= conj.(batched_results_reverse[a][b])`
	`225`	`+ end`
`223`	`226`	`copyto!(`
`224`	`227`	`view(compressed_matrix_reverse, 1 + ((a - 1) * Br + (b - 1)) % Nr, :),`
`225`	`228`	`vec(batched_results_reverse[a][b]),`