add chain function

Author: simeonschaub

src/ChainRulesCore.jl

@@ -4,8 +4,8 @@ using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broad
 export frule, rrule
 export wirtinger_conjugate, wirtinger_primal, refine_differential
 export @scalar_rule, @thunk
-export extern, cast, store!
-export Wirtinger, Zero, One, Casted, DNE, Thunk, InplaceableThunk
+export extern, chain, cast, store!
+export Wirtinger, ComplexGradient, Zero, One, Casted, DNE, Thunk, InplaceableThunk
 export NO_FIELDS
 
 include("differentials.jl")

src/differential_arithmetic.jl

@@ -114,3 +114,40 @@ for T in (:Any,)
     @eval Base.:*(a::AbstractThunk, b::$T) = extern(a) * b
     @eval Base.:*(a::$T, b::AbstractThunk) = a * extern(b)
 end
+
+function chain(outer, inner; swap_order=false)
+    if swap_order
+        return Wirtinger(
+                         wirtinger_primal(inner) * wirtinger_primal(outer) +
+                         conj(wirtinger_conjugate(inner)) * wirtinger_conjugate(outer),
+                         wirtinger_conjugate(inner) * wirtinger_primal(outer) +
+                         conj(wirtinger_primal(inner) * wirtinger_conjugate(outer))
+                        ) |> refine_differential
+    end
+    return Wirtinger(
+                     wirtinger_primal(outer) * wirtinger_primal(inner) +
+                     wirtinger_conjugate(outer) * conj(wirtinger_conjugate(inner)),
+                     wirtinger_primal(outer) * wirtinger_conjugate(inner) +
+                     wirtinger_conjugate(outer) * conj(wirtinger_primal(inner))
+                    ) |> refine_differential
+end
+
+function chain(outer::ComplexGradient, inner; swap_order=false)
+    if swap_order
+        return ComplexGradient(
+                               wirtinger_conjugate(inner) + conj(wirtinger_primal(inner)) * 
+                               outer.val
+                              )
+    end
+    return ComplexGradient(
+                           outer.val *
+                           (wirtinger_conjugate(inner) + conj(wirtinger_primal(inner)))
+                          )
+end
+
+function chain(outer::ComplexGradient, inner::ComplexGradient; swap_order=false)
+    if swap_order
+        return ComplexGradient(conj(inner.val) * outer.val)
+    end
+    return ComplexGradient(outer.val * conj(inner.val))
+end

src/differentials.jl

@@ -307,14 +307,18 @@ function itself, when that function is not a closure.
 const NO_FIELDS = DNE()
 
 """
-    refine_differential(𝒟::Type, der)
+    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`.
 """
 function refine_differential(::Type{<:Union{<:Real, AbstractArray{<:Real}}}, w::Wirtinger)
+    w = refine_differential(w)
     return wirtinger_primal(w) + wirtinger_conjugate(w)
 end
-refine_differential(::Any, der) = der  # most of the time leave it alone.
+refine_differential(::Any, der) = refine_differential(der)  # most of the time leave it alone.
+
+refine_differential(w::Wirtinger{<:Any,Zero}) = w.primal
+refine_differential(der::Any) = der