WIP work on new @scalar_rule

make real scalar rules work.

correct @scalarrule forward rule return

Wirtinger scalar working

work WirtingerRule test as a test of @scalar_rule

Fix spelling

Co-Authored-By: simeonschaub <simeondavidschaub99@gmail.com>
Oxford Comma

Co-Authored-By: simeonschaub <simeondavidschaub99@gmail.com>
spelling

Co-Authored-By: Nick Robinson <npr251@gmail.com>
docstring for propagator_name
spelling

Co-Authored-By: Nick Robinson <npr251@gmail.com>

Author: oxinabox

src/ChainRulesCore.jl

@@ -2,6 +2,7 @@ module ChainRulesCore
 using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broadcastable
 
 export AbstractRule, Rule, frule, rrule
+export wirtinger_conjugate, wirtinger_primal, differential
 export @scalar_rule, @thunk
 export extern, cast, store!
 export Wirtinger, Zero, One, Casted, DNE, Thunk, DNERule
@@ -11,5 +12,5 @@ include("differentials.jl")
 include("differential_arithmetic.jl")
 include("rule_types.jl")
 include("rules.jl")
-#include("rule_definition_tools.jl")
+include("rule_definition_tools.jl")
 end # module

src/differentials.jl

@@ -230,8 +230,22 @@ Base.show(io::IO, x::Thunk) = println(io, "Thunk($(repr(x.f)))")
 
 """
     NO_FIELDS
+
 Constant for the reverse-mode derivative with respect to a structure that has no fields.
 The most notable use for this is for the reverse-mode derivative with respect to the
 function itself, when that function is not a closure.
 """
 const NO_FIELDS = DNE()
+
+####
+"""
+    differential(𝒟::Type, der)
+
+For some differential (e.g. a `Number`, `AbstractDifferential`, `Matrix`, etc.),
+convert it to another differential that is more suited for the domain given by
+the type 𝒟.
+"""
+function differential(::Type{<:Union{<:Real, AbstractArray{<:Real}}}, w::Wirtinger)
+    return wirtinger_primal(w) + wirtinger_conjugate(w)
+end
+differential(::Any, der) = der  # most of the time leave it alone.

src/rule_definition_tools.jl

@@ -1,5 +1,89 @@
 # These are some macros (and supporting functions) to make it easier to define rules.
 
+"""
+    propagator_name(f, propname)
+Determines a reasonable name for the propagator function.
+The name doesn't really matter too much as it is a local function to be returned
+by `frule` or `rrule`, but a good name make debugging easier.
+`f` should be some form of AST representation of the actual function,
+`propname` should be either `:pullback` or `:pushforward`
+
+This is able to deal with fairly complex expressions for `f`:
+
+    julia> propagator_name(:bar, :pushforward)
+    :bar_pushforward
+
+    julia> propagator_name(esc(:(Base.Random.foo)), :pullback)
+    :foo_pullback
+"""
+propagator_name(f::Expr, propname::Symbol) = propagator_name(f.args[end], propname)
+propagator_name(fname::Symbol, propname::Symbol) = Symbol(fname, :_, propname)
+propagator_name(fname::QuoteNode, propname::Symbol) = propagator_name(fname.value, propname)
+
+
+"""
+    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)
+    # This is basically Δs ⋅ ∂s
+
+    # Notice: the thunking of `∂s[i] (potentially) saves us some computation
+    # if `Δs[i]` is a `AbstractDifferential` otherwise it is computed as soon
+    # as the pullback is evaluated
+    ∂_mul_Δs = [:(@thunk($(∂s[i])) * $(Δs[i])) for i in 1:length(∂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)
+        differential($𝒟, w)
+    end
+end
+
 """
     @scalar_rule(f(x₁, x₂, ...),
                  @setup(statement₁, statement₂, ...),
@@ -42,7 +126,7 @@ e.g. `f(x₁::Complex, x₂)`, which will constrain `x₁` to `Complex` and `x
 At present this does not support defining for closures/functors.
 Thus in reverse-mode, the first returned partial,
 representing the derivative with respect to the function itself, is always `NO_FIELDS`.
-And in forwards-mode, the first input to the returned propergator is always ignored.
+And in forward-mode, the first input to the returned propagator is always ignored.
 
 The result of `f(x₁, x₂, ...)` is automatically bound to `Ω`. This
 allows the primal result to be conveniently referenced (as `Ω`) within the
@@ -69,6 +153,9 @@ For examples, see ChainRulesCore' `rules` directory.
 See also: [`frule`](@ref), [`rrule`](@ref), [`AbstractRule`](@ref)
 """
 macro scalar_rule(call, maybe_setup, partials...)
+    ############################################################################
+    # Setup: normalizing input form etc
+
     if Meta.isexpr(maybe_setup, :macrocall) && maybe_setup.args[1] == Symbol("@setup")
         setup_stmts = map(esc, maybe_setup.args[3:end])
     else
@@ -77,6 +164,7 @@ macro scalar_rule(call, maybe_setup, partials...)
     end
     @assert Meta.isexpr(call, :call)
     f = esc(call.args[1])
+
     # Annotate all arguments in the signature as scalars
     inputs = map(call.args[2:end]) do arg
         esc(Meta.isexpr(arg, :(::)) ? arg : Expr(:(::), arg, :Number))
@@ -90,6 +178,7 @@ macro scalar_rule(call, maybe_setup, partials...)
         end
     end
 
+    # For consistency in code that follows we make all partials tuple expressions
     partials = map(partials) do partial
         if Meta.isexpr(partial, :tuple)
             partial
@@ -98,59 +187,58 @@ macro scalar_rule(call, maybe_setup, partials...)
             Expr(:tuple, partial)
         end
     end
-    @show partials
-
-    ############################################################
-    # Make pullback
-    #(TODO: move to own function)
-    # TODO: Wirtinger
-    
-    Δs = [Symbol(string(:Δ, i)) for i in 1:length(partials)]
-    pullback_returns = map(eachindex(inputs)) do input_i
-        ∂s = [partials.args[input_i] for partial in partials]
-        ∂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
-        # as the pullback is evaluated
-        ∂_mul_Δs = [:(@thunk($(∂s[i])) * $(Δs[i])) for i in 1:length(∂s)]
-        :(+($(∂_mul_Δs...)))
-    else
 
-    pullback = quote
-        function $(Symbol(nameof(f), :_pullback))($(Δs...))
-            return (ChainRulesCore.NO_FIELDS, $(pullback_returns...))
+    ############################################################################
+    # Main body: defining the results of the frule/rrule
+
+    # An expression that when evaluated will return the type of the input domain.
+    # Multiple repetitions of this expression should optimize ot. 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]...)))))
+
+    n_outputs = length(partials)
+    n_inputs = length(inputs)
+
+    pushforward = let
+        # Δs is the input to the propagator rule
+        # because this is push-forward there is one per input to the function
+        Δ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)
         end
-    end
 
-    ########################################
-    quote
-        function ChainRulesCore.rrule(::typeof($f), $(inputs...))
-            $(esc(:Ω)) = $call
-            $(setup_stmts...)
-            return $(esc(:Ω)), $esc(pullback)
+        quote
+            # _ is the input derivative w.r.t. function internals. since we do not
+            # allow closures/functors with @scalar_rule, it is always ignored
+            function $(propagator_name(f, :pushforward))(_, $(Δs...))
+                return $(Expr(:tuple, pushforward_returns...))
+            end
         end
     end
-end
-#==
-    if !all(Meta.isexpr(partial, :tuple) 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 = map(1:length(inputs) do i
-            reverse_partials = [partial.args[i] for partial in partials]
-            push!(reverse_rules, rule_from_partials(inputs[i], reverse_partials...))
+
+    pullback = let
+        # Δs is the input to the propagator rule
+        # because this is a pull-back there is one per output of function
+        Δs = [Symbol(string(:Δ, i)) for i in 1:n_outputs]
+
+        # 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)
+        end
+
+        quote
+            function $(propagator_name(f, :pullback))($(Δs...))
+                return (NO_FIELDS, $(pullback_returns...))
+            end
         end
-    else
-        @assert length(inputs) == 1 && all(!Meta.isexpr(partial, :tuple) for partial in partials)
-        forward_rules = Any[rule_from_partials(inputs[1], partial) for partial in partials]
-        reverse_rules = Any[rule_from_partials(inputs[1], partials...)]
     end
 
-    # First pseudo-partial is derivative WRT function itself.  Since this macro does not
-    # support closures, it is just the empty NamedTuple
-    forward_rules = Expr(:tuple, ZERO_RULE, forward_rules...)
-    reverse_rules = Expr(:tuple, NO_FIELDS, reverse_rules...)
-    return quote
+    ############################################################################
+    # Final return: building the expression to insert in the place of this macro
+
+    code = quote
         if fieldcount(typeof($f)) > 0
             throw(ArgumentError(
                 "@scalar_rule cannot be used on closures/functors (such as $f)"
@@ -160,57 +248,13 @@ end
         function ChainRulesCore.frule(::typeof($f), $(inputs...))
             $(esc(:Ω)) = $call
             $(setup_stmts...)
-            return $(esc(:Ω)), $forward_rules
+            return $(esc(:Ω)), $pushforward
         end
+
         function ChainRulesCore.rrule(::typeof($f), $(inputs...))
             $(esc(:Ω)) = $call
             $(setup_stmts...)
-            return $(esc(:Ω)), $reverse_rules
-        end
-    end
-end
-==#
-
-@macroexpand(@scalar_rule(one(x), Zero()))
-
-
-
-#==
-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)]
-    Δs_tuple = Expr(:tuple, Δs...)
-    if isempty(wirtinger_indices)
-        ∂_mul_Δs = [:(@thunk($(∂s[i])) * $(Δs[i])) for i in 1:length(∂s)]
-        return :(Rule($Δs_tuple -> +($(∂_mul_Δs...))))
-    else
-        ∂_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_rule = :(Rule($Δs_tuple -> +($(∂_mul_Δs_primal...))))
-        conjugate_rule = :(Rule($Δs_tuple -> +($(∂_mul_Δs_conjugate...))))
-        return quote
-            $(∂_wirtinger_defs...)
-            AbstractRule(typeof($input_arg), $primal_rule, $conjugate_rule)
+            return $(esc(:Ω)), $pullback
         end
     end
 end
-==#

test/differentials.jl

@@ -79,4 +79,20 @@
         ]
         @test isempty(ambig_methods)
     end
+
+
+    @testset "Differential" begin
+        @test differential(typeof(1.0 + 1im), Wirtinger(2,2)) == Wirtinger(2,2)
+        @test differential(typeof([1.0 + 1im]), Wirtinger(2,2)) == Wirtinger(2,2)
+
+        @test differential(typeof(1.2), Wirtinger(2,2)) == 4
+        @test differential(typeof([1.2]), Wirtinger(2,2)) == 4
+
+        # For most differentials, in most domains, this does nothing
+        for der in (DNE(), @thunk(23), @thunk(Wirtinger(2,2)), [1 2], One(), Zero(), 0.0)
+            for 𝒟 in typeof.((1.0 + 1im, [1.0 + 1im], 1.2, [1.2]))
+                @test differential(𝒟, der) === der
+            end
+        end
+    end
 end

test/rule_types.jl

@@ -1,8 +1,8 @@
 
 @testset "rule types" begin
-    #==
+    # The following is deprecated and should be remove next release
     @testset "iterating and indexing rules" begin
-        _, rule = frule(dummy_identity, 1)
+        rule = Rule(identity)
         i = 0
         for r in rule
             @test r === rule
@@ -12,8 +12,8 @@
         @test rule[1] == rule
         @test_throws BoundsError rule[2]
     end
-    ==#
-    
+
+
     @testset "Rule" begin
         @testset "show" begin
             @test occursin(r"^Rule\(.*foo.*\)$", repr(Rule(function foo() 1 end)))