Merge pull request #79 from JuliaDiff/ox/nowirtinger

Remove Wirtinger

Author: oxinabox

src/ChainRulesCore.jl

@@ -2,16 +2,14 @@ module ChainRulesCore
 using Base.Broadcast: broadcasted, Broadcasted, broadcastable, materialize, materialize!
 
 export frule, rrule
-export refine_differential, wirtinger_conjugate, wirtinger_primal
 export @scalar_rule, @thunk
 export extern, store!, unthunk
-export Composite, DoesNotExist, InplaceableThunk, One, Thunk, Wirtinger, Zero
+export Composite, DoesNotExist, InplaceableThunk, One, Thunk, Zero
 export NO_FIELDS
 
 include("compat.jl")
 
 include("differentials/abstract_differential.jl")
-include("differentials/wirtinger.jl")
 include("differentials/zero.jl")
 include("differentials/does_not_exist.jl")
 include("differentials/one.jl")

src/differential_arithmetic.jl

@@ -8,44 +8,12 @@ Thus we can avoid any ambiguities.
 
 Notice:
     The precedence goes:
-    `Wirtinger, Zero, DoesNotExist, One, AbstractThunk, Composite, Any`
+    `Zero, DoesNotExist, One, AbstractThunk, Composite, Any`
     Thus each of the @eval loops creating definitions of + and *
     defines the combination this type with all types of  lower precidence.
     This means each eval loops is 1 item smaller than the previous.
 ==#
 
-
-function Base.:*(a::Wirtinger, b::Wirtinger)
-    error("""
-          Cannot multiply two Wirtinger objects; this error likely means a
-          `WirtingerRule` was inappropriately defined somewhere. Multiplication
-          of two Wirtinger objects is not defined because chain rule application
-          often expands into a non-commutative operation in the Wirtinger
-          calculus. To put it another way: simply given two Wirtinger objects
-          and no other information, we can't know "locally" which components to
-          conjugate in order to implement the chain rule. We could pick a
-          convention; for example, we could define `a::Wirtinger * b::Wirtinger`
-          such that we assume the chain rule application is of the form `f_a ∘ f_b`
-          instead of `f_b ∘ f_a`. However, picking such a convention is likely to
-          lead to silently incorrect derivatives due to commutativity assumptions
-          in downstream generic code that deals with the reals. Thus, ChainRulesCore
-          makes this operation an error instead.
-          """)
-end
-
-function Base.:+(a::Wirtinger, b::Wirtinger)
-    return Wirtinger(+(a.primal, b.primal), a.conjugate + b.conjugate)
-end
-
-for T in (:Zero, :DoesNotExist, :One, :AbstractThunk, :Any)
-    @eval Base.:+(a::Wirtinger, b::$T) = a + Wirtinger(b, Zero())
-    @eval Base.:+(a::$T, b::Wirtinger) = Wirtinger(a, Zero()) + b
-
-    @eval Base.:*(a::Wirtinger, b::$T) = Wirtinger(a.primal * b, a.conjugate * b)
-    @eval Base.:*(a::$T, b::Wirtinger) = Wirtinger(a * b.primal, a * b.conjugate)
-end
-
-
 Base.:+(::Zero, b::Zero) = Zero()
 Base.:*(::Zero, ::Zero) = Zero()
 for T in (:DoesNotExist, :One, :AbstractThunk, :Any)

src/differentials/abstract_differential.jl

@@ -40,13 +40,3 @@ wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
 @inline extern(x) = x
 
 @inline Base.conj(x::AbstractDifferential) = x
-
-"""
-    refine_differential(𝒟::Type, der)
-
-Converts, if required, a differential object `der`
-(e.g. a `Number`, `AbstractDifferential`, `Matrix`, etc.),
-to another  differential that is more suited for the domain given by the type 𝒟.
-Often this will behave as the identity function on `der`.
-"""
-refine_differential(::Any, der) = der  # most of the time leave it alone.

src/differentials/wirtinger.jl

@@ -74,14 +74,8 @@ macro scalar_rule(call, maybe_setup, partials...)
     )
     f = call.args[1]
 
-    # An expression that when evaluated will return the type of the input domain.
-    # Multiple repetitions of this expression should optimize out. But if it does not then
-    # may need to move its definition into the body of the `rrule`/`frule`
-    𝒟 = :(typeof(first(promote($(call.args[2:end]...)))))
-
-    frule_expr = scalar_frule_expr(𝒟, f, call, setup_stmts, inputs, partials)
-    rrule_expr = scalar_rrule_expr(𝒟, f, call, setup_stmts, inputs, partials)
-
+    frule_expr = scalar_frule_expr(f, call, setup_stmts, inputs, partials)
+    rrule_expr = scalar_rrule_expr(f, call, setup_stmts, inputs, partials)
 
     ############################################################################
     # Final return: building the expression to insert in the place of this macro
@@ -147,7 +141,7 @@ function _normalize_scalarrules_macro_input(call, maybe_setup, partials)
     return call, setup_stmts, inputs, partials
 end
 
-function scalar_frule_expr(𝒟, f, call, setup_stmts, inputs, partials)
+function scalar_frule_expr(f, call, setup_stmts, inputs, partials)
     n_outputs = length(partials)
     n_inputs = length(inputs)
 
@@ -156,7 +150,7 @@ function scalar_frule_expr(𝒟, f, call, setup_stmts, inputs, partials)
     Δs = [Symbol(string(:Δ, i)) for i in 1:n_inputs]
     pushforward_returns = map(1:n_outputs) do output_i
         ∂s = partials[output_i].args
-        propagation_expr(𝒟, Δs, ∂s)
+        propagation_expr(Δs, ∂s)
     end
     if n_outputs > 1
         # For forward-mode we only return a tuple if output actually a tuple.
@@ -182,7 +176,7 @@ function scalar_frule_expr(𝒟, f, call, setup_stmts, inputs, partials)
     end
 end
 
-function scalar_rrule_expr(𝒟, f, call, setup_stmts, inputs, partials)
+function scalar_rrule_expr(f, call, setup_stmts, inputs, partials)
     n_outputs = length(partials)
     n_inputs = length(inputs)
 
@@ -193,7 +187,7 @@ function scalar_rrule_expr(𝒟, f, call, setup_stmts, inputs, partials)
     # 1 partial derivative per input
     pullback_returns = map(1:n_inputs) do input_i
         ∂s = [partial.args[input_i] for partial in partials]
-        propagation_expr(𝒟, Δs, ∂s)
+        propagation_expr(Δs, ∂s)
     end
 
     pullback = quote
@@ -212,30 +206,14 @@ function scalar_rrule_expr(𝒟, f, call, setup_stmts, inputs, partials)
 end
 
 """
-    propagation_expr(𝒟, Δs, ∂s)
+    propagation_expr(Δs, ∂s)
 
     Returns the expression for the propagation of
     the input gradient `Δs` though the partials `∂s`.
-
-    𝒟 is an expression that when evaluated returns the type-of the input domain.
-    For example if the derivative is being taken at the point `1` it returns `Int`.
-    if it is taken at `1+1im` it returns `Complex{Int}`.
-    At present it is ignored for non-Wirtinger derivatives.
 """
-function propagation_expr(𝒟, Δs, ∂s)
-    wirtinger_indices = findall(∂s) do ex
-        Meta.isexpr(ex, :call) && ex.args[1] === :Wirtinger
-    end
-    ∂s = map(esc, ∂s)
-    if isempty(wirtinger_indices)
-        return standard_propagation_expr(Δs, ∂s)
-    else
-        return wirtinger_propagation_expr(𝒟, wirtinger_indices, Δs, ∂s)
-    end
-end
-
-function standard_propagation_expr(Δs, ∂s)
+function propagation_expr(Δs, ∂s)
     # This is basically Δs ⋅ ∂s
+    ∂s = map(esc, ∂s)
 
     # Notice: the thunking of `∂s[i] (potentially) saves us some computation
     # if `Δs[i]` is a `AbstractDifferential` otherwise it is computed as soon
@@ -244,36 +222,6 @@ function standard_propagation_expr(Δs, ∂s)
     return :(+($(∂_mul_Δs...)))
 end
 
-function wirtinger_propagation_expr(𝒟, wirtinger_indices, Δs, ∂s)
-    ∂_mul_Δs_primal = Any[]
-    ∂_mul_Δs_conjugate = Any[]
-    ∂_wirtinger_defs = Any[]
-    for i in 1:length(∂s)
-        if i in wirtinger_indices
-            Δi = Δs[i]
-            ∂i = Symbol(string(:∂, i))
-            push!(∂_wirtinger_defs, :($∂i = $(∂s[i])))
-            ∂f∂i_mul_Δ = :(wirtinger_primal($∂i) * wirtinger_primal($Δi))
-            ∂f∂ī_mul_Δ̄ = :(conj(wirtinger_conjugate($∂i)) * wirtinger_conjugate($Δi))
-            ∂f̄∂i_mul_Δ = :(wirtinger_conjugate($∂i) * wirtinger_primal($Δi))
-            ∂f̄∂ī_mul_Δ̄ = :(conj(wirtinger_primal($∂i)) * wirtinger_conjugate($Δi))
-            push!(∂_mul_Δs_primal, :($∂f∂i_mul_Δ + $∂f∂ī_mul_Δ̄))
-            push!(∂_mul_Δs_conjugate, :($∂f̄∂i_mul_Δ + $∂f̄∂ī_mul_Δ̄))
-        else
-            ∂_mul_Δ = :(@thunk($(∂s[i])) * $(Δs[i]))
-            push!(∂_mul_Δs_primal, ∂_mul_Δ)
-            push!(∂_mul_Δs_conjugate, ∂_mul_Δ)
-        end
-    end
-    primal_sum = :(+($(∂_mul_Δs_primal...)))
-    conjugate_sum = :(+($(∂_mul_Δs_conjugate...)))
-    return quote  # This will be a block, so will have value equal to last statement
-        $(∂_wirtinger_defs...)
-        w = Wirtinger($primal_sum, $conjugate_sum)
-        refine_differential($𝒟, w)
-    end
-end
-
 """
     propagator_name(f, propname)
 

src/rule_definition_tools.jl

@@ -38,114 +38,3 @@ _second(t) = Base.tuple_type_head(Base.tuple_type_tail(t))
     @test rrx == 2
     @test rr1 == 1
 end
-
-
-@testset "Basic Wirtinger scalar_rule" begin
-    myabs2(x) = abs2(x)
-    @scalar_rule(myabs2(x), Wirtinger(x', x))
-
-    @testset "real input" begin
-        # even though our rule was define in terms of Wirtinger,
-        # pushforward result will be real as real (even if seed is Compex)
-
-        x = rand(Float64)
-        f, myabs2_pushforward = frule(myabs2, x)
-        @test f === x^2
-
-        Δ = One()
-        df = @inferred myabs2_pushforward(NamedTuple(), Δ)
-        @test df === x + x
-
-        Δ = rand(Complex{Int64})
-        df = @inferred myabs2_pushforward(NamedTuple(), Δ)
-        @test df === Δ * (x + x)
-    end
-
-    @testset "complex input" begin
-        z = rand(Complex{Float64})
-        f, myabs2_pushforward = frule(myabs2, z)
-        @test f === abs2(z)
-
-        df = @inferred myabs2_pushforward(NamedTuple(), One())
-        @test df === Wirtinger(z', z)
-
-        Δ = rand(Complex{Int64})
-        df = @inferred myabs2_pushforward(NamedTuple(), Δ)
-        @test df === Wirtinger(Δ * z', Δ * z)
-    end
-end
-
-
-@testset "Advanced Wirtinger @scalar_rule: abs_to_pow" begin
-    # This is based on SimeonSchaub excellent example:
-    # https://gist.github.com/simeonschaub/a6dfcd71336d863b3777093b3b8d9c97
-
-    # This is much more complex than the previous case
-    # as it has many different types
-    # depending on input, and the output types do not always agree
-
-    abs_to_pow(x, p) = abs(x)^p
-    @scalar_rule(
-        abs_to_pow(x::Real, p),
-        (
-            p == 0 ? Zero() : p * abs_to_pow(x, p-1) * sign(x),
-            Ω * log(abs(x))
-        )
-    )
-
-    @scalar_rule(
-        abs_to_pow(x::Complex, p),
-        @setup(u = abs(x)),
-        (
-            p == 0 ? Zero() : p * u^(p-1) * Wirtinger(x' / 2u, x / 2u),
-            Ω * log(abs(x))
-        )
-    )
-
-
-    f = abs_to_pow
-    @testset "f($x, $p)" for (x, p) in Iterators.product(
-        (2, 3.4, -2.1, -10+0im, 2.3-2im),
-        (0, 1, 2, 4.3, -2.1, 1+.2im)
-    )
-        expected_type_df_dx =
-            if iszero(p)
-                Zero
-            elseif typeof(x) <: Complex
-                Wirtinger
-            elseif typeof(p) <: Complex
-                Complex
-            else
-                Real
-            end
-
-        expected_type_df_dp =
-            if typeof(p) <: Real
-                Real
-            else
-                Complex
-            end
-
-
-        res = frule(f, x, p)
-        @test res !== nothing  # Check the rule was defined
-        fx, f_pushforward = res
-        df(Δx, Δp) = f_pushforward(NamedTuple(), Δx, Δp)
-
-        df_dx::Thunk = df(One(), Zero())
-        df_dp::Thunk = df(Zero(), One())
-        @test fx == f(x, p)  # Check we still get the normal value, right
-        @test df_dx() isa expected_type_df_dx
-        @test df_dp() isa expected_type_df_dp
-
-
-        res = rrule(f, x, p)
-        @test res !== nothing  # Check the rule was defined
-        fx, f_pullback = res
-        dself, df_dx, df_dp = f_pullback(One())
-        @test fx == f(x, p)  # Check we still get the normal value, right
-        @test dself == NO_FIELDS
-        @test df_dx() isa expected_type_df_dx
-        @test df_dp() isa expected_type_df_dp
-    end
-end