Get rid of all Rule types, add InplacableThunk

spelling is hard

Author: oxinabox

src/ChainRulesCore.jl

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

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, :Any)
+    The precidence goes: (:Wirtinger, :Casted, :Zero, :DNE, :One, :Thunk, :InplaceableThunk, :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, :Any)
+for T in (:Casted, :Zero, :DNE, :One, :Thunk, :InplaceableThunk, :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, :Any)
+for T in (:Zero, :DNE, :One, :Thunk, :InplaceableThunk, :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, :Any)
+for T in (:DNE, :One, :Thunk, :InplaceableThunk, :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, :Any)
+for T in (:One, :Thunk, :InplaceableThunk, :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, :Any)
+for T in (:Thunk, :InplaceableThunk, :Any)
     @eval Base.:+(a::One, b::$T) = extern(a) + b
     @eval Base.:+(a::$T, b::One) = a + extern(b)
 
@@ -91,10 +91,21 @@ end
 
 Base.:+(a::Thunk, b::Thunk) = extern(a) + extern(b)
 Base.:*(a::Thunk, b::Thunk) = extern(a) * extern(b)
-for T in (:Any,)  #This loop is redundant but for consistency...
+for T in (:InplaceableThunk, :Any)
     @eval Base.:+(a::Thunk, b::$T) = extern(a) + b
     @eval Base.:+(a::$T, b::Thunk) = 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)
+end

src/differentials.jl

@@ -246,6 +246,43 @@ Base.conj(x::Thunk) = @thunk(conj(extern(x)))
 
 Base.show(io::IO, x::Thunk) = println(io, "Thunk($(repr(x.f)))")
 
+"""
+    InplaceableThunk(val::Thunk, add!::Function)
+
+A wrapper for a `Thunk`, that allows it to define an inplace `add!` function,
+which is used internally in `accumulate!(Δ, ::InplaceableThunk)`.
+
+`add!` should be defined such that: `ithunk.add!(Δ) = Δ .+= ithunk.val`
+but it should do this more efficently than simply doing this directly.
+(Otherwise one can just use a normal `Thunk`).
+
+Most operations on an `InplaceableThunk` treat it just like a normal `Thunk`;
+and destroy its inplacability.
+"""
+struct InplaceableThunk{T<:Thunk, F} <: AbstractDifferential
+    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!(Δ)
+
+
 """
     NO_FIELDS
 
@@ -255,7 +292,6 @@ function itself, when that function is not a closure.
 """
 const NO_FIELDS = DNE()
 
-####
 """
     differential(𝒟::Type, der)
 

src/operations.jl

@@ -0,0 +1,45 @@
+# TODO: This all needs a fair bit of rethinking
+
+"""
+    accumulate(Δ, ∂)
+
+Return `Δ + ∂` evaluated in a manner that supports ChainRulesCore's
+various `AbstractDifferential` types.
+
+See also: [`accumulate!`](@ref), [`store!`](@ref), [`AbstractRule`](@ref)
+"""
+accumulate(Δ, ∂) = Δ .+ ∂
+
+"""
+    accumulate!(Δ, ∂)
+
+Similar to [`accumulate`](@ref), but attempts to compute `Δ + rule(args...)` in-place,
+storing the result in `Δ`.
+
+Note: this function may not actually store the result in `Δ` if `Δ` is immutable,
+so it is best to always call this as `Δ = accumulate!(Δ, ∂)` just in-case.
+
+This function is overloadable by [`InplaceableThunk`s](@ref).
+See also: [`accumulate`](@ref), [`store!`](@ref), [`AbstractRule`](@ref)
+"""
+function accumulate!(Δ, ∂)
+    return materialize!(Δ, broadcastable(cast(Δ) + ∂))
+end
+
+accumulate!(Δ::Number, ∂) = accumulate(Δ, ∂)
+
+
+
+"""
+    store!(Δ, ∂)
+
+Stores `∂`, in `Δ`, overwriting what ever was in `Δ` before.
+potentially avoiding intermediate temporary allocations that might be
+necessary for alternative approaches  (e.g. `copyto!(Δ, extern(∂))`)
+
+Like [`accumulate`](@ref) and [`accumulate!`](@ref), this function is intended
+to be customizable for specific rules/input types.
+
+See also: [`accumulate`](@ref), [`accumulate!`](@ref), [`AbstractRule`](@ref)
+"""
+store!(Δ, ∂) = materialize!(Δ, broadcastable(∂))