fix: complex support for wrong-mode pushforward/pullback (#733)

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.42"
+version = "0.6.43"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/src/first_order/pullback.jl

@@ -151,26 +151,58 @@ function _pullback_via_pushforward(
     f::F,
     pushforward_prep::PushforwardPrep,
     backend::AbstractADType,
-    x::Number,
+    x::Real,
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    t1 = pushforward(f, pushforward_prep, backend, x, (one(x),), contexts...)
-    dx = dot(only(t1), dy)
+    a = only(pushforward(f, pushforward_prep, backend, x, (one(x),), contexts...))
+    dx = dot(a, dy)
     return dx
 end
 
 function _pullback_via_pushforward(
     f::F,
     pushforward_prep::PushforwardPrep,
     backend::AbstractADType,
-    x::AbstractArray,
+    x::Complex,
+    dy,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    a = only(pushforward(f, pushforward_prep, backend, x, (one(x),), contexts...))
+    b = only(pushforward(f, pushforward_prep, backend, x, (im * one(x),), contexts...))
+    dx = real(dot(a, dy)) + im * real(dot(b, dy))
+    return dx
+end
+
+function _pullback_via_pushforward(
+    f::F,
+    pushforward_prep::PushforwardPrep,
+    backend::AbstractADType,
+    x::AbstractArray{<:Real},
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
     dx = map(CartesianIndices(x)) do j
-        t1 = pushforward(f, pushforward_prep, backend, x, (basis(x, j),), contexts...)
-        dot(only(t1), dy)
+        a = only(pushforward(f, pushforward_prep, backend, x, (basis(x, j),), contexts...))
+        dot(a, dy)
+    end
+    return dx
+end
+
+function _pullback_via_pushforward(
+    f::F,
+    pushforward_prep::PushforwardPrep,
+    backend::AbstractADType,
+    x::AbstractArray{<:Complex},
+    dy,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    dx = map(CartesianIndices(x)) do j
+        a = only(pushforward(f, pushforward_prep, backend, x, (basis(x, j),), contexts...))
+        b = only(
+            pushforward(f, pushforward_prep, backend, x, (im * basis(x, j),), contexts...),
+        )
+        real(dot(a, dy)) + im * real(dot(b, dy))
     end
     return dx
 end
@@ -236,12 +268,43 @@ function _pullback_via_pushforward(
     y,
     pushforward_prep::PushforwardPrep,
     backend::AbstractADType,
-    x::Number,
+    x::Real,
+    dy,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    a = only(pushforward(f!, y, pushforward_prep, backend, x, (one(x),), contexts...))
+    dx = dot(a, dy)
+    return dx
+end
+
+function _pullback_via_pushforward(
+    f!::F,
+    y,
+    pushforward_prep::PushforwardPrep,
+    backend::AbstractADType,
+    x::Complex,
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    t1 = pushforward(f!, y, pushforward_prep, backend, x, (one(x),), contexts...)
-    dx = dot(only(t1), dy)
+    a = only(pushforward(f!, y, pushforward_prep, backend, x, (one(x),), contexts...))
+    b = only(pushforward(f!, y, pushforward_prep, backend, x, (im * one(x),), contexts...))
+    dx = real(dot(a, dy)) + im * real(dot(b, dy))
+    return dx
+end
+
+function _pullback_via_pushforward(
+    f!::F,
+    y,
+    pushforward_prep::PushforwardPrep,
+    backend::AbstractADType,
+    x::AbstractArray{<:Real},
+    dy,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    dx = map(CartesianIndices(x)) do j  # preserve shape
+        a = only(pushforward(f!, y, pushforward_prep, backend, x, (basis(x, j),), contexts...))
+        dot(a, dy)
+    end
     return dx
 end
 
@@ -250,13 +313,18 @@ function _pullback_via_pushforward(
     y,
     pushforward_prep::PushforwardPrep,
     backend::AbstractADType,
-    x::AbstractArray,
+    x::AbstractArray{<:Complex},
     dy,
     contexts::Vararg{Context,C},
 ) where {F,C}
     dx = map(CartesianIndices(x)) do j  # preserve shape
-        t1 = pushforward(f!, y, pushforward_prep, backend, x, (basis(x, j),), contexts...)
-        dot(only(t1), dy)
+        a = only(pushforward(f!, y, pushforward_prep, backend, x, (basis(x, j),), contexts...))
+        b = only(
+            pushforward(
+                f!, y, pushforward_prep, backend, x, (im * basis(x, j),), contexts...
+            ),
+        )
+        real(dot(a, dy)) + im * real(dot(b, dy))
     end
     return dx
 end

DifferentiationInterface/src/first_order/pushforward.jl

@@ -157,13 +157,28 @@ function _pushforward_via_pullback(
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
-    t1 = pullback(f, pullback_prep, backend, x, (one(y),), contexts...)
-    dy = dot(only(t1), dx)
+    a = only(pullback(f, pullback_prep, backend, x, (one(y),), contexts...))
+    dy = dot(a, dx)
     return dy
 end
 
 function _pushforward_via_pullback(
-    y::AbstractArray,
+    y::Complex,
+    f::F,
+    pullback_prep::PullbackPrep,
+    backend::AbstractADType,
+    x,
+    dx,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    a = only(pullback(f, pullback_prep, backend, x, (one(y),), contexts...))
+    b = only(pullback(f, pullback_prep, backend, x, (im * one(y),), contexts...))
+    dy = real(dot(a, dx)) + im * real(dot(b, dx))
+    return dy
+end
+
+function _pushforward_via_pullback(
+    y::AbstractArray{<:Real},
     f::F,
     pullback_prep::PullbackPrep,
     backend::AbstractADType,
@@ -172,8 +187,25 @@ function _pushforward_via_pullback(
     contexts::Vararg{Context,C},
 ) where {F,C}
     dy = map(CartesianIndices(y)) do i
-        t1 = pullback(f, pullback_prep, backend, x, (basis(y, i),), contexts...)
-        dot(only(t1), dx)
+        a = only(pullback(f, pullback_prep, backend, x, (basis(y, i),), contexts...))
+        dot(a, dx)
+    end
+    return dy
+end
+
+function _pushforward_via_pullback(
+    y::AbstractArray{<:Complex},
+    f::F,
+    pullback_prep::PullbackPrep,
+    backend::AbstractADType,
+    x,
+    dx,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    dy = map(CartesianIndices(y)) do i
+        a = only(pullback(f, pullback_prep, backend, x, (basis(y, i),), contexts...))
+        b = only(pullback(f, pullback_prep, backend, x, (im * basis(y, i),), contexts...))
+        real(dot(a, dx)) + im * real(dot(b, dx))
     end
     return dy
 end
@@ -236,16 +268,35 @@ end
 
 function _pushforward_via_pullback(
     f!::F,
-    y::AbstractArray,
+    y::AbstractArray{<:Real},
+    pullback_prep::PullbackPrep,
+    backend::AbstractADType,
+    x,
+    dx,
+    contexts::Vararg{Context,C},
+) where {F,C}
+    dy = map(CartesianIndices(y)) do i  # preserve shape
+        a = only(pullback(f!, y, pullback_prep, backend, x, (basis(y, i),), contexts...))
+        dot(a, dx)
+    end
+    return dy
+end
+
+function _pushforward_via_pullback(
+    f!::F,
+    y::AbstractArray{<:Complex},
     pullback_prep::PullbackPrep,
     backend::AbstractADType,
     x,
     dx,
     contexts::Vararg{Context,C},
 ) where {F,C}
     dy = map(CartesianIndices(y)) do i  # preserve shape
-        t1 = pullback(f!, y, pullback_prep, backend, x, (basis(y, i),), contexts...)
-        dot(only(t1), dx)
+        a = only(pullback(f!, y, pullback_prep, backend, x, (basis(y, i),), contexts...))
+        b = only(
+            pullback(f!, y, pullback_prep, backend, x, (im * basis(y, i),), contexts...)
+        )
+        real(dot(a, dx)) + im * real(dot(b, dx))
     end
     return dy
 end

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.9.4"
+version = "0.9.5"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterfaceTest/src/scenarios/complex.jl

@@ -1,3 +1,33 @@
+square_only(x::AbstractVector) = only(x)^2
+abs2_only(x::AbstractVector) = abs2(only(x))
+
+function complex_holomorphic_gradient_scenarios()
+    # http://arxiv.org/abs/2409.06752
+    dy = 1.0
+    x = [1.0 + im]
+    grad = 2 * conj(x)
+    scens = Scenario[
+        Scenario{:gradient,:out}(square_only, x; res1=grad),
+        Scenario{:gradient,:in}(square_only, x; res1=grad),
+        Scenario{:pullback,:out}(square_only, x; tang=(dy,), res1=(grad,)),
+        Scenario{:pullback,:in}(square_only, x; tang=(dy,), res1=(grad,)),
+    ]
+    return scens
+end
+
+function complex_gradient_scenarios()
+    dy = 1.0
+    x = [1.0 + im]
+    grad = 2 * x
+    scens = Scenario[
+        Scenario{:gradient,:out}(abs2_only, x; res1=grad),
+        Scenario{:gradient,:in}(abs2_only, x; res1=grad),
+        Scenario{:pullback,:out}(abs2_only, x; tang=(dy,), res1=(grad,)),
+        Scenario{:pullback,:in}(abs2_only, x; tang=(dy,), res1=(grad,)),
+    ]
+    return scens
+end
+
 """
     complex_scenarios()
 
@@ -15,8 +45,6 @@ function complex_scenarios()
     dy_6 = float.(-5:2:5) .+ im
     dy_12 = float.(-11:2:11) .+ im
 
-    V = Vector{Complex{Float64}}
-
     scens = vcat(
         # one argument
         num_to_num_scenarios(x_; dx=dx_, dy=dy_),
@@ -26,6 +54,9 @@ function complex_scenarios()
         # two arguments
         num_to_vec_scenarios_twoarg(x_; dx=dx_, dy=dy_6),
         vec_to_vec_scenarios_twoarg(x_6; dx=dx_6, dy=dy_12),
+        # complex gradients
+        complex_gradient_scenarios(),
+        complex_holomorphic_gradient_scenarios(),
     )
 
     return filter(s -> !(operator(s) in SECOND_ORDER), scens)