Support complex output arrays in derivative(f!, y, x)

src/apiutils.jl

@@ -40,6 +40,16 @@ end
     return Expr(:tuple, [:(single_seed(Partials{N,V}, Val{$i}())) for i in 1:N]...)
 end
 
+# A buffer for a complex-valued output stores `Complex{Dual}`, not a dual over a complex value: real
+# and imaginary parts each carry their own partials, matching what `dual_definition_retval` returns
+# for complex-valued primitives. `V` is the underlying real scalar type in both cases.
+const DualBuffer{T,V,N} = Union{AbstractArray{Dual{T,V,N}},AbstractArray{Complex{Dual{T,V,N}}}}
+
+# Builds a buffer element carrying `v` as its value and `partials` on every real component.
+@inline buffer_element(::Type{D}, v, partials) where {D<:Dual} = D(v, partials)
+@inline buffer_element(::Type{Complex{D}}, v, partials) where {D<:Dual} =
+    Complex(D(real(v), partials), D(imag(v), partials))
+
 # Only seed indices that are structurally non-zero
 structural_eachindex(x::AbstractArray) = structural_eachindex(x, x)
 function structural_eachindex(x::AbstractArray, y::AbstractArray)
@@ -73,29 +83,30 @@ end
 # Copies the values of `x` into `duals` with zero partials. Used both to remove seeds `duals` is
 # currently carrying and to initialize a freshly allocated work buffer, whose elements must all be
 # written before the target function reads them.
-seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x) where {T,V,N} =
+seed_zero_partials!(duals::DualBuffer{T,V,N}, x) where {T,V,N} =
     _seed_zero_partials!(duals, x, structural_eachindex(duals, x))
 
 # Zeroes the partials of `count` elements starting at structural position `index`. Chunk mode only
 # needs to clear the chunk it just seeded, so writing through to the end of the array would be O(n)
 # redundant work per chunk, i.e. O(n^2/N) per sweep. `count` mirrors the `chunksize` argument of
 # `seed!(duals, x, index, seeds, chunksize)`.
-function seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, index,
+function seed_zero_partials!(duals::DualBuffer{T,V,N}, x, index,
                              count = N) where {T,V,N}
     idxs = Iterators.take(Iterators.drop(structural_eachindex(duals, x), index - 1), count)
     return _seed_zero_partials!(duals, x, idxs)
 end
 
-function _seed_zero_partials!(duals::AbstractArray{Dual{T,V,N}}, x, idxs) where {T,V,N}
+function _seed_zero_partials!(duals::DualBuffer{T,V,N}, x, idxs) where {T,V,N}
     seed = zero(Partials{N,V})
+    E = eltype(duals)
     if isbitstype(V)
         for idx in idxs
-            duals[idx] = Dual{T,V,N}(x[idx], seed)
+            duals[idx] = buffer_element(E, x[idx], seed)
         end
     else
         for idx in idxs
             if isassigned(x, idx)
-                duals[idx] = Dual{T,V,N}(x[idx], seed)
+                duals[idx] = buffer_element(E, x[idx], seed)
             else
                 Base._unsetindex!(duals, idx)
             end

src/config.jl

@@ -63,6 +63,11 @@ struct DerivativeConfig{T,D} <: AbstractConfig{1}
     duals::D
 end
 
+# Work-buffer element type for an output whose scalar type is `Y`; see `DualBuffer`.
+@inline dual_buffer_eltype(::Type{T}, ::Type{Y}, ::Val{N}) where {T,Y<:Real,N} = Dual{T,Y,N}
+@inline dual_buffer_eltype(::Type{T}, ::Type{Y}, ::Val{N}) where {T,Y<:Complex,N} =
+    Complex{Dual{T,real(Y),N}}
+
 """
     ForwardDiff.DerivativeConfig(f!, y::AbstractArray, x::Real)
 
@@ -82,8 +87,8 @@ This constructor does not store/modify `y` or `x`.
 function DerivativeConfig(f::F,
                           y::AbstractArray{Y},
                           x::X,
-                          tag::T = Tag(f, X)) where {F,X<:Real,Y<:Real,T}
-    duals = similar(y, Dual{T,Y,1})
+                          tag::T = Tag(f, X)) where {F,X<:Real,Y<:Union{Real,Complex},T}
+    duals = similar(y, dual_buffer_eltype(T, Y, Val(1)))
     return DerivativeConfig{T,typeof(duals)}(duals)
 end
 

src/dual.jl

@@ -98,6 +98,10 @@ Dual{T,V,N}(x::Base.TwicePrecision) where {T,V,N} =
     end
 end
 
+# A complex value carries its dual information componentwise (see `buffer_element`).
+@inline value(z::Complex) = complex(value(real(z)), value(imag(z)))
+@inline value(::Type{T}, z::Complex) where {T} = complex(value(T, real(z)), value(T, imag(z)))
+
 @inline partials(x) = Partials{0,typeof(x)}(tuple())
 @inline partials(d::Dual) = d.partials
 @inline partials(x, i...) = zero(x)

test/AllocationsTest.jl

@@ -32,6 +32,15 @@ convert_test_574() = convert(ForwardDiff.Dual{Nothing,ForwardDiff.Dual{Nothing,F
     allocs_convert_test_574() = @allocated convert_test_574()
     allocs_convert_test_574()
     @test iszero(allocs_convert_test_574())
+
+    # same, for a `Complex{Dual}` buffer: `_seed_zero_partials!` reads the element type off `duals`,
+    # which must constant-fold
+    yc = Vector{ComplexF64}(undef, 1000)
+    cduals = ForwardDiff.DerivativeConfig(nothing, yc, 0.0).duals
+    allocs_szp!(cduals, yc)
+    @test iszero(allocs_szp!(cduals, yc))
+    allocs_szp!(cduals, yc, 1, 1)
+    @test iszero(allocs_szp!(cduals, yc, 1, 1))
 end
 
 @testset "Test jacobian! allocations" begin

test/DerivativeTest.jl

@@ -110,8 +110,71 @@ end
     @test_throws DimensionMismatch ForwardDiff.derivative(sum, fill(2pi, 3))
 end
 
+# `abs`/`conj`/`real`/`angle` are nowhere complex differentiable, and work only because the real and
+# imaginary parts carry separate partials.
+const COMPLEX_OUTPUT_FUNCS = (
+    ("cis and exp",    (y, x) -> (y[1] = cis(x); y[2] = exp((1+2im)*x)),
+                       x -> [cis(x), exp((1+2im)*x)]),
+    ("abs and conj",   (y, x) -> (z = cis(x)*(2+x); y[1] = abs(z)+0im; y[2] = conj(z)),
+                       x -> (z = cis(x)*(2+x); [abs(z)+0im, conj(z)])),
+    ("real and angle", (y, x) -> (z = (1+2im)*x^2+3; y[1] = real(z)+0im; y[2] = angle(z)+0im),
+                       x -> (z = (1+2im)*x^2+3; [real(z)+0im, angle(z)+0im])),
+    ("sqrt and log",   (y, x) -> (z = 2cis(x)+3; y[1] = sqrt(z); y[2] = log(z)/(x+1im)),
+                       x -> (z = 2cis(x)+3; [sqrt(z), log(z)/(x+1im)])),
+)
+
 @testset "complex output" begin
     @test ForwardDiff.derivative(x -> (1+im)*x, 0) == (1+im)
+
+    # The in-place path must agree exactly with the non-mutating one: same `Complex{Dual}` arithmetic.
+    # `y` only matches approximately, since recomputing `f(x)` in `Float64` reassociates.
+    @testset "in-place, $name" for (name, f!, f) in COMPLEX_OUTPUT_FUNCS
+        x = 0.7
+        v, d = f(x), ForwardDiff.derivative(f, x)
+        @test !(eltype(d) <: ForwardDiff.Dual)
+        @test d ≈ Calculus.derivative(f, x) atol=FINITEDIFF_ERROR
+
+        y = Vector{ComplexF64}(undef, 2)
+        for cfg in ((), (ForwardDiff.DerivativeConfig(f!, y, x),))
+            @test ForwardDiff.derivative(f!, y, x, cfg...) == d
+            @test y ≈ v
+
+            out = similar(d)
+            @test ForwardDiff.derivative!(out, f!, y, x, cfg...) === out
+            @test out == d
+
+            out = DiffResults.DiffResult(similar(v), similar(d))
+            @test ForwardDiff.derivative!(out, f!, y, x, cfg...) === out
+            @test DiffResults.value(out) ≈ v
+            @test DiffResults.derivative(out) == d
+        end
+    end
+
+    @testset "in-place, entries f! leaves alone" begin
+        y = ComplexF64[0, 9-4im]
+        d = ForwardDiff.derivative((y, x) -> (y[1] = cis(x)), y, 0.3)
+        @test d[1] ≈ im*cis(0.3)
+        @test d[2] === 0.0+0.0im
+        @test y[2] === 9.0-4.0im
+    end
+
+    @testset "in-place, nested" begin
+        h!(y, x) = (y[1] = cis(2x); y[2] = x^2 + 3im*x)
+        @test ForwardDiff.derivative(1.0) do x
+            ForwardDiff.derivative(h!, Vector{Complex{typeof(x)}}(undef, 2), x)
+        end ≈ [-4cis(2.0), 2.0+0im]
+    end
+
+    # `Complex{BigFloat}` is not `isbitstype`, covering the `isassigned` branch of
+    # `_seed_zero_partials!`; the `Matrix` covers a non-vector shape.
+    @testset "in-place, $(eltype(y))" for y in (Vector{Complex{BigFloat}}(undef, 3),
+                                               Matrix{ComplexF64}(undef, 2, 2))
+        f(x) = fill(cis(x)*(1+x), size(y))
+        x = convert(real(eltype(y)), 4//10)
+        @test ForwardDiff.derivative((y, x) -> (y .= cis(x)*(1+x)), y, x) ==
+              ForwardDiff.derivative(f, x)
+        @test y ≈ f(x)
+    end
 end
 
 @testset "NaN-safe mode" begin