introduce a function unwrap_wirtinger

Author: simeonschaub

src/differential_arithmetic.jl

@@ -115,7 +115,10 @@ for T in (:Any,)
     @eval Base.:*(a::$T, b::AbstractThunk) = a * extern(b)
 end
 
-@inline function chain(outer, inner, swap_order=false)
+@inline chain(outer, inner, swap_order=false) =
+    _chain(unwrap_wirtiner(outer), unwrap_wirtinger(inner), swap_order)
+
+@inline function _chain(outer, inner, swap_order)
     if swap_order
         return Wirtinger(
                          wirtinger_primal(inner) * wirtinger_primal(outer) +
@@ -132,7 +135,7 @@ end
                     ) |> refine_differential
 end
 
-@inline function chain(outer::ComplexGradient, inner, swap_order=false)
+@inline function _chain(outer::ComplexGradient, inner, swap_order)
     if swap_order
         return ComplexGradient(
                                (wirtinger_conjugate(inner) + conj(wirtinger_primal(inner))) * 
@@ -145,7 +148,7 @@ end
                           )
 end
 
-@inline function chain(outer::ComplexGradient, inner::ComplexGradient, swap_order=false)
+@inline function _chain(outer::ComplexGradient, inner::ComplexGradient, swap_order)
     if swap_order
         return ComplexGradient(conj(inner.val) * outer.val)
     end

src/differentials.jl

@@ -47,8 +47,15 @@ wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
 
 abstract type AbstractWirtinger <: AbstractDifferential end
 
+unwrap_wirtinger(x) = x
+unwrap_wirtinger(x::Union{Casted,AbstractThunk}) = unwrap_wirtinger(extern(x))
+
 wirtinger_primal(x) = x
+wirtinger_primal(x::Union{Casted,AbstractThunk}) =
+    throw(ArgumentError("`wirtinger_primal` is not defined for $(typeof(x)). Call `unwrap_wirtinger` first")
 wirtinger_conjugate(::Any) = Zero()
+wirtinger_primal(x::Union{Casted,AbstractThunk}) =
+    throw(ArgumentError("`wirtinger_conjugate` is not defined for $(typeof(x)). Call `unwrap_wirtinger` first")
 
 extern(x::AbstractWirtinger) = throw(ArgumentError("`AbstractWirtinger` cannot be converted to an external type."))