much AbstractThunk

Author: oxinabox

src/differential_arithmetic.jl

@@ -7,7 +7,7 @@ subtypes, as we know the full set that might be encountered.
 Thus we can avoid any ambiguities.
 
 Notice:
-    The precidence goes: (:Wirtinger, :Casted, :Zero, :DNE, :One, :Thunk, :InplaceableThunk, :Any)
+    The precidence goes: (:Wirtinger, :Casted, :Zero, :DNE, :One, :AbstractThunk, :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.
@@ -36,7 +36,7 @@ function Base.:+(a::Wirtinger, b::Wirtinger)
     return Wirtinger(+(a.primal, b.primal), a.conjugate + b.conjugate)
 end
 
-for T in (:Casted, :Zero, :DNE, :One, :Thunk, :InplaceableThunk, :Any)
+for T in (:Casted, :Zero, :DNE, :One, :AbstractThunk, :Any)
     @eval Base.:+(a::Wirtinger, b::$T) = a + Wirtinger(b, Zero())
     @eval Base.:+(a::$T, b::Wirtinger) = Wirtinger(a, Zero()) + b
 
@@ -47,7 +47,7 @@ end
 
 Base.:+(a::Casted, b::Casted) = Casted(broadcasted(+, a.value, b.value))
 Base.:*(a::Casted, b::Casted) = Casted(broadcasted(*, a.value, b.value))
-for T in (:Zero, :DNE, :One, :Thunk, :InplaceableThunk, :Any)
+for T in (:Zero, :DNE, :One, :AbstractThunk, :Any)
     @eval Base.:+(a::Casted, b::$T) = Casted(broadcasted(+, a.value, b))
     @eval Base.:+(a::$T, b::Casted) = Casted(broadcasted(+, a, b.value))
 
@@ -58,7 +58,7 @@ end
 
 Base.:+(::Zero, b::Zero) = Zero()
 Base.:*(::Zero, ::Zero) = Zero()
-for T in (:DNE, :One, :Thunk, :InplaceableThunk, :Any)
+for T in (:DNE, :One, :AbstractThunk, :Any)
     @eval Base.:+(::Zero, b::$T) = b
     @eval Base.:+(a::$T, ::Zero) = a
 
@@ -69,7 +69,7 @@ end
 
 Base.:+(::DNE, ::DNE) = DNE()
 Base.:*(::DNE, ::DNE) = DNE()
-for T in (:One, :Thunk, :InplaceableThunk, :Any)
+for T in (:One, :AbstractThunk, :Any)
     @eval Base.:+(::DNE, b::$T) = b
     @eval Base.:+(a::$T, ::DNE) = a
 
@@ -80,7 +80,7 @@ end
 
 Base.:+(a::One, b::One) = extern(a) + extern(b)
 Base.:*(::One, ::One) = One()
-for T in (:Thunk, :InplaceableThunk, :Any)
+for T in (:AbstractThunk, :Any)
     @eval Base.:+(a::One, b::$T) = extern(a) + b
     @eval Base.:+(a::$T, b::One) = a + extern(b)
 
@@ -89,23 +89,12 @@ for T in (:Thunk, :InplaceableThunk, :Any)
 end
 
 
-Base.:+(a::Thunk, b::Thunk) = extern(a) + extern(b)
-Base.:*(a::Thunk, b::Thunk) = extern(a) * extern(b)
-for T in (:InplaceableThunk, :Any)
-    @eval Base.:+(a::Thunk, b::$T) = extern(a) + b
-    @eval Base.:+(a::$T, b::Thunk) = a + extern(b)
+Base.:+(a::AbstractThunk, b::AbstractThunk) = extern(a) + extern(b)
+Base.:*(a::AbstractThunk, b::AbstractThunk) = extern(a) * extern(b)
+for T in (:Any,)
+    @eval Base.:+(a::AbstractThunk, b::$T) = extern(a) + b
+    @eval Base.:+(a::$T, b::AbstractThunk) = a + extern(b)
 
-    @eval Base.:*(a::Thunk, b::$T) = extern(a) * b
-    @eval Base.:*(a::$T, b::Thunk) = a * extern(b)
-end
-
-# InplaceableThunk acts just like Thunk
-Base.:+(a::InplaceableThunk, b::InplaceableThunk) = extern(a) + extern(b)
-Base.:*(a::InplaceableThunk, b::InplaceableThunk) = extern(a) * extern(b)
-for T in (:Any, )
-    @eval Base.:+(a::InplaceableThunk, b::$T) = extern(a) + b
-    @eval Base.:+(a::$T, b::InplaceableThunk) = a + extern(b)
-
-    @eval Base.:*(a::InplaceableThunk, b::$T) = extern(a) * b
-    @eval Base.:*(a::$T, b::InplaceableThunk) = a * extern(b)
+    @eval Base.:*(a::AbstractThunk, b::$T) = extern(a) * b
+    @eval Base.:*(a::$T, b::AbstractThunk) = a * extern(b)
 end

src/differentials.jl

@@ -173,6 +173,24 @@ Base.iterate(x::One) = (x, nothing)
 Base.iterate(::One, ::Any) = nothing
 
 
+#####
+##### `AbstractThunk
+#####
+abstract type AbstractThunk <: AbstractDifferential end
+
+Base.Broadcast.broadcastable(x::AbstractThunk) = broadcastable(extern(x))
+
+@inline function Base.iterate(x::AbstractThunk)
+    externed = extern(x)
+    element, state = iterate(externed)
+    return element, (externed, state)
+end
+
+@inline function Base.iterate(::AbstractThunk, (externed, state))
+    element, new_state = iterate(externed, state)
+    return element, (externed, new_state)
+end
+
 #####
 ##### `Thunk`
 #####
@@ -218,31 +236,20 @@ itself a `Thunk`.
 If you got the expression from another `rrule` (or `frule`), you don't need to
 `@thunk` it since it will have been thunked if required, by the defining rule.
 """
-struct Thunk{F} <: AbstractDifferential
+struct Thunk{F} <: AbstractThunk
     f::F
 end
 
 macro thunk(body)
     return :(Thunk(() -> $(esc(body))))
 end
 
-(x::Thunk)() = x.f()
-@inline extern(x::Thunk) = extern(x())
-
-Base.Broadcast.broadcastable(x::Thunk) = broadcastable(extern(x))
+# have to define this here after `@thunk` and `Thunk` is defined
+Base.conj(x::AbstractThunk) = @thunk(conj(extern(x)))
 
-@inline function Base.iterate(x::Thunk)
-    externed = extern(x)
-    element, state = iterate(externed)
-    return element, (externed, state)
-end
 
-@inline function Base.iterate(::Thunk, (externed, state))
-    element, new_state = iterate(externed, state)
-    return element, (externed, new_state)
-end
-
-Base.conj(x::Thunk) = @thunk(conj(extern(x)))
+(x::Thunk)() = x.f()
+@inline extern(x::Thunk) = extern(x())
 
 Base.show(io::IO, x::Thunk) = println(io, "Thunk($(repr(x.f)))")
 
@@ -259,29 +266,21 @@ but it should do this more efficently than simply doing this directly.
 Most operations on an `InplaceableThunk` treat it just like a normal `Thunk`;
 and destroy its inplacability.
 """
-struct InplaceableThunk{T<:Thunk, F} <: AbstractDifferential
+struct InplaceableThunk{T<:Thunk, F} <: AbstractThunk
     val::T
     add!::F
 end
 
 (x::InplaceableThunk)() = x.val()
 @inline extern(x::InplaceableThunk) = extern(x.val)
 
-Base.Broadcast.broadcastable(x::InplaceableThunk) = broadcastable(x.val)
-
-@inline function Base.iterate(x::InplaceableThunk, args...)
-    return iterate(x.val, args...)
-end
-
-Base.conj(x::InplaceableThunk) = conj(x.val)
-
 function Base.show(io::IO, x::InplaceableThunk)
     println(io, "InplaceableThunk($(repr(x.val)), $(repr(x.add!)))")
 end
 
 # The real reason we have this:
 accumulate!(Δ, ∂::InplaceableThunk) = ∂.add!(Δ)
-store!(Δ, ∂::InplaceableThunk) = ∂.add!((Δ.*=false))  # zero i, then add to it.
+store!(Δ, ∂::InplaceableThunk) = ∂.add!((Δ.*=false))  # zero it, then add to it.
 
 """
     NO_FIELDS