Merge pull request #26 from JuliaDiff/oc/revert13

Revert #13

Author: oxinabox

src/differentials.jl

@@ -303,3 +303,26 @@ 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,16 @@ 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(input, partial.args...) for (input, partial) in zip(inputs, partials)]
+        forward_rules = Any[rule_from_partials(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(inputs[i], reverse_partials...))
+            push!(reverse_rules, rule_from_partials(reverse_partials...))
         end
     else
         @assert length(inputs) == 1 && all(!Meta.isexpr(partial, :tuple) for partial in partials)
-        forward_rules = Any[rule_from_partials(input, partial) for (input, partial) in zip(inputs, partials)]
-        reverse_rules = Any[rule_from_partials(inputs[1], partials...)]
+        forward_rules = Any[rule_from_partials(partial) for partial in partials]
+        reverse_rules = Any[rule_from_partials(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 +107,7 @@ macro scalar_rule(call, maybe_setup, partials...)
     end
 end
 
-function rule_from_partials(input_arg, ∂s...)
+function rule_from_partials(∂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 +140,7 @@ function rule_from_partials(input_arg, ∂s...)
         conjugate_rule = :(Rule($Δs_tuple -> add($(∂_mul_Δs_conjugate...))))
         return quote
             $(∂_wirtinger_defs...)
-            WirtingerRule(typeof($input_arg), $primal_rule, $conjugate_rule)
+            WirtingerRule($primal_rule, $conjugate_rule)
         end
     end
 end

src/rules.jl

@@ -184,25 +184,13 @@ DNERule(args...) = DNE()
 #####
 
 """
-    WirtingerRule([𝒟::Type, ]P::AbstractRule, C::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 `𝒟` is specified, return a `Rule` evaluating
-to `P(Δ) + C(Δ)` if `𝒟 <: Real`, otherwise return `WirtingerRule(P, C)`.
+TODO
 """
 struct WirtingerRule{P<:AbstractRule,C<:AbstractRule} <: AbstractRule
     primal::P
     conjugate::C
 end
 
-function WirtingerRule(𝒟::Type, primal::AbstractRule, conjugate::AbstractRule)
-    if 𝒟 <: Real || eltype(𝒟) <: Real
-        return Rule((args...) -> add(primal(args...), conjugate(args...)))
-    else
-        return WirtingerRule(primal, conjugate)
-    end
-end
-
 function (rule::WirtingerRule)(args...)
     return Wirtinger(rule.primal(args...), rule.conjugate(args...))
 end

test/rules.jl

@@ -42,41 +42,4 @@ 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), WirtingerRule(
-                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, WirtingerRule
+    DNE, Thunk, Casted, DNERule
 using Base.Broadcast: broadcastable
 
 @testset "ChainRulesCore" begin