fix wirtinger mechanism (#29)

* fix wirtinger mechanism

* another attempt...

* implement proposed changes

* fix tests accordingly

* fix another omission

Author: simeonschaub

src/differentials.jl

@@ -303,26 +303,3 @@ add_thunk(a, b::Thunk) = add(a, extern(b))
 mul_thunk(a::Thunk, b::Thunk) = mul(extern(a), extern(b))
 mul_thunk(a::Thunk, b) = mul(extern(a), b)
 mul_thunk(a, b::Thunk) = mul(a, extern(b))
-
-#####
-##### misc.
-#####
-
-"""
-    Wirtinger(primal::Real, conjugate::Real)
-
-Return `add(primal, conjugate)`.
-
-Actually implementing the Wirtinger calculus generally requires that the
-summed terms of the Wirtinger differential (`∂f/∂z * dz` and `∂f/∂z̄ * dz̄`) be
-stored individually. However, if both of these terms are real-valued, then
-downstream Wirtinger propagation mechanisms resolve to the same mechanisms as
-real-valued calculus, so that the terms' sum can be eagerly computed and
-propagated without requiring a special `Wirtinger` representation
-
-This method primarily exists as an optimization.
-"""
-function Wirtinger(primal::Union{Real,DNE,Zero,One},
-                   conjugate::Union{Real,DNE,Zero,One})
-    return add(primal, conjugate)
-end

src/rule_definition_tools.jl

@@ -80,16 +80,17 @@ macro scalar_rule(call, maybe_setup, partials...)
         end
     end
     if all(Meta.isexpr(partial, :tuple) for partial in partials)
-        forward_rules = Any[rule_from_partials(partial.args...) for partial in partials]
+        input_rep = :(first(promote($(inputs...))))  # stand-in with the right type for an input
+        forward_rules = Any[rule_from_partials(input_rep, partial.args...) for partial in partials]
         reverse_rules = Any[]
         for i in 1:length(inputs)
             reverse_partials = [partial.args[i] for partial in partials]
-            push!(reverse_rules, rule_from_partials(reverse_partials...))
+            push!(reverse_rules, rule_from_partials(inputs[i], reverse_partials...))
         end
     else
         @assert length(inputs) == 1 && all(!Meta.isexpr(partial, :tuple) for partial in partials)
-        forward_rules = Any[rule_from_partials(partial) for partial in partials]
-        reverse_rules = Any[rule_from_partials(partials...)]
+        forward_rules = Any[rule_from_partials(inputs[1], partial) for partial in partials]
+        reverse_rules = Any[rule_from_partials(inputs[1], partials...)]
     end
     forward_rules = length(forward_rules) == 1 ? forward_rules[1] : Expr(:tuple, forward_rules...)
     reverse_rules = length(reverse_rules) == 1 ? reverse_rules[1] : Expr(:tuple, reverse_rules...)
@@ -107,7 +108,7 @@ macro scalar_rule(call, maybe_setup, partials...)
     end
 end
 
-function rule_from_partials(∂s...)
+function rule_from_partials(input_arg, ∂s...)
     wirtinger_indices = findall(x -> Meta.isexpr(x, :call) && x.args[1] === :Wirtinger,  ∂s)
     ∂s = map(esc, ∂s)
     Δs = [Symbol(string(:Δ, i)) for i in 1:length(∂s)]
@@ -140,7 +141,7 @@ function rule_from_partials(∂s...)
         conjugate_rule = :(Rule($Δs_tuple -> add($(∂_mul_Δs_conjugate...))))
         return quote
             $(∂_wirtinger_defs...)
-            WirtingerRule($primal_rule, $conjugate_rule)
+            AbstractRule(typeof($input_arg), $primal_rule, $conjugate_rule)
         end
     end
 end

src/rules.jl

@@ -184,7 +184,13 @@ DNERule(args...) = DNE()
 #####
 
 """
-TODO
+    WirtingerRule(primal::AbstractRule, conjugate::AbstractRule)
+
+Construct a `WirtingerRule` object, which is an `AbstractRule` that consists of
+an `AbstractRule` for both the primal derivative ``∂/∂x`` and the conjugate
+derivative ``∂/∂x̅``. If the domain `𝒟` of the function might be real, consider
+calling `AbstractRule(𝒟, primal, conjugate)` instead, to make use of a more
+efficient representation wherever possible.
 """
 struct WirtingerRule{P<:AbstractRule,C<:AbstractRule} <: AbstractRule
     primal::P
@@ -195,6 +201,20 @@ function (rule::WirtingerRule)(args...)
     return Wirtinger(rule.primal(args...), rule.conjugate(args...))
 end
 
+"""
+    AbstractRule(𝒟::Type, primal::AbstractRule, conjugate::AbstractRule)
+
+Return a `Rule` evaluating to `primal(Δ) + conjugate(Δ)` if `𝒟 <: Real`,
+otherwise return `WirtingerRule(P, C)`.
+"""
+function AbstractRule(𝒟::Type, primal::AbstractRule, conjugate::AbstractRule)
+    if 𝒟 <: Real || eltype(𝒟) <: Real
+        return Rule((args...) -> add(primal(args...), conjugate(args...)))
+    else
+        return WirtingerRule(primal, conjugate)
+    end
+end
+
 #####
 ##### `frule`/`rrule`
 #####

test/rules.jl

@@ -42,4 +42,41 @@ dummy_identity(x) = x
         @test rule[1] == rule
         @test_throws BoundsError rule[2]
     end
+
+    @testset "WirtingerRule" begin
+        myabs2(x) = abs2(x)
+
+        function frule(::typeof(myabs2), x)
+            return abs2(x), AbstractRule(
+                typeof(x),
+                Rule(Δx -> Δx * x'),
+                Rule(Δx -> Δx * x)
+            )
+        end
+
+        # real input
+        x = rand(Float64)
+        f, _df = @inferred frule(myabs2, x)
+        @test f === x^2
+
+        df = @inferred _df(One())
+        @test df === x + x
+
+        Δ = rand(Complex{Int64})
+        df = @inferred _df(Δ)
+        @test df === Δ * (x + x)
+
+
+        # complex input
+        z = rand(Complex{Float64})
+        f, _df = @inferred frule(myabs2, z)
+        @test f === abs2(z)
+
+        df = @inferred _df(One())
+        @test df === Wirtinger(z', z)
+
+        Δ = rand(Complex{Int64})
+        df = @inferred _df(Δ)
+        @test df === Wirtinger(Δ * z', Δ * z)
+    end
 end

test/runtests.jl

@@ -5,7 +5,7 @@ using LinearAlgebra: Diagonal
 using ChainRulesCore: extern, accumulate, accumulate!, store!, @scalar_rule,
     Wirtinger, wirtinger_primal, wirtinger_conjugate, add_wirtinger, mul_wirtinger,
     Zero, add_zero, mul_zero, One, add_one, mul_one, Casted, cast, add_casted, mul_casted,
-    DNE, Thunk, Casted, DNERule
+    DNE, Thunk, Casted, DNERule, WirtingerRule
 using Base.Broadcast: broadcastable
 
 @testset "ChainRulesCore" begin