Remove all uses of Cassette and define the arithmetic directly (#35)

* =Remove all uses of Cassette and difine the arithmetic directly

* =Remove use of add from accumulate and define single-arg +

* Update test/rule_types.jl

* Update src/differential_arithmetic.jl

* Update test/rules.jl

* clean up setup code

* use infix

* make ambiguity test not count as 20,000 tests

* remove whitespace

* fix and confirm fixed typos in Wirtinger arthmetic

Author: oxinabox

Project.toml

@@ -1,12 +1,8 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "0.1.1"
-
-[deps]
-Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
+version = "0.2.0"
 
 [compat]
-Cassette = "^0.2"
 julia = "^1.0"
 
 [extras]

src/ChainRulesCore.jl

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

src/differential_arithmetic.jl

@@ -0,0 +1,100 @@
+#==
+All differentials need to define + and *.
+That happens here.
+
+We just use @eval to define all the combinations for AbstractDifferential
+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)
+    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.
+==#
+
+
+function Base.:*(a::Wirtinger, b::Wirtinger)
+    error("""
+          Cannot multiply two Wirtinger objects; this error likely means a
+          `WirtingerRule` was inappropriately defined somewhere. Multiplication
+          of two Wirtinger objects is not defined because chain rule application
+          often expands into a non-commutative operation in the Wirtinger
+          calculus. To put it another way: simply given two Wirtinger objects
+          and no other information, we can't know "locally" which components to
+          conjugate in order to implement the chain rule. We could pick a
+          convention; for example, we could define `a::Wirtinger * b::Wirtinger`
+          such that we assume the chain rule application is of the form `f_a ∘ f_b`
+          instead of `f_b ∘ f_a`. However, picking such a convention is likely to
+          lead to silently incorrect derivatives due to commutativity assumptions
+          in downstream generic code that deals with the reals. Thus, ChainRulesCore
+          makes this operation an error instead.
+          """)
+end
+
+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)
+    @eval Base.:+(a::Wirtinger, b::$T) = a + Wirtinger(b, Zero())
+    @eval Base.:+(a::$T, b::Wirtinger) = Wirtinger(a, Zero()) + b
+
+    @eval Base.:*(a::Wirtinger, b::$T) = Wirtinger(a.primal * b, a.conjugate * b)
+    @eval Base.:*(a::$T, b::Wirtinger) = Wirtinger(a * b.primal, a * b.conjugate)
+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)
+    @eval Base.:+(a::Casted, b::$T) = Casted(broadcasted(+, a.value, b))
+    @eval Base.:+(a::$T, b::Casted) = Casted(broadcasted(+, a, b.value))
+
+    @eval Base.:*(a::Casted, b::$T) = Casted(broadcasted(*, a.value, b))
+    @eval Base.:*(a::$T, b::Casted) = Casted(broadcasted(*, a, b.value))
+end
+
+
+Base.:+(::Zero, b::Zero) = Zero()
+Base.:*(::Zero, ::Zero) = Zero()
+for T in (:DNE, :One, :Thunk, :Any)
+    @eval Base.:+(::Zero, b::$T) = b
+    @eval Base.:+(a::$T, ::Zero) = a
+
+    @eval Base.:*(::Zero, ::$T) = Zero()
+    @eval Base.:*(::$T, ::Zero) = Zero()
+end
+
+
+Base.:+(::DNE, ::DNE) = DNE()
+Base.:*(::DNE, ::DNE) = DNE()
+for T in (:One, :Thunk, :Any)
+    @eval Base.:+(::DNE, b::$T) = b
+    @eval Base.:+(a::$T, ::DNE) = a
+
+    @eval Base.:*(::DNE, ::$T) = DNE()
+    @eval Base.:*(::$T, ::DNE) = DNE()
+end
+
+
+Base.:+(a::One, b::One) = extern(a) + extern(b)
+Base.:*(::One, ::One) = One()
+for T in (:Thunk, :Any)
+    @eval Base.:+(a::One, b::$T) = extern(a) + b
+    @eval Base.:+(a::$T, b::One) = a + extern(b)
+
+    @eval Base.:*(::One, b::$T) = b
+    @eval Base.:*(a::$T, ::One) = a
+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...
+    @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

src/differentials.jl

@@ -9,9 +9,9 @@ support, broadcast fusion, zero-elision, etc. into nicely separated parts.
 
 All subtypes of `AbstractDifferential` implement the following operations:
 
-`add(a, b)`: linearly combine differential `a` and differential `b`
+`+(a, b)`: linearly combine differential `a` and differential `b`
 
-`mul(a, b)`: multiply the differential `a` by the differential `b`
+`*(a, b)`: multiply the differential `a` by the differential `b`
 
 `Base.conj(x)`: complex conjugate of the differential `x`
 
@@ -26,6 +26,8 @@ Additionally, all subtypes of `AbstractDifferential` support `Base.iterate` and
 """
 abstract type AbstractDifferential end
 
+Base.:+(x::AbstractDifferential) = x
+
 """
     extern(x)
 
@@ -39,40 +41,6 @@ wrapped by `x`, such that mutating `extern(x)` might mutate `x` itself.
 
 @inline Base.conj(x::AbstractDifferential) = x
 
-#=
-This `AbstractDifferential` algebra has a monad-y "fallthrough" implementation;
-each step handles an element of the algebra before dispatching to the next step.
-This way, we don't need to implement promotion/conversion rules between subtypes
-of `AbstractDifferential` to resolve potential ambiguities.
-=#
-
-const PRECEDENCE_LIST = [:wirtinger, :casted, :zero, :dne, :one, :thunk, :fallback]
-
-global defs = Expr(:block)
-
-let previous_add_name = :add, previous_mul_name = :mul
-    for name in PRECEDENCE_LIST
-        next_add_name = Symbol(string(:add_, name))
-        next_mul_name = Symbol(string(:mul_, name))
-        push!(defs.args, quote
-            @inline $(previous_add_name)(a, b) = $(next_add_name)(a, b)
-            @inline $(previous_mul_name)(a, b) = $(next_mul_name)(a, b)
-        end)
-        previous_add_name = next_add_name
-        previous_mul_name = next_mul_name
-    end
-end
-
-eval(defs)
-
-@inline add_fallback(a, b) = a + b
-
-@inline mul_fallback(a, b) = a * b
-
-@inline add(x) = x
-
-@inline mul(x) = x
-
 #####
 ##### `Wirtinger`
 #####
@@ -120,33 +88,6 @@ Base.iterate(::Wirtinger, ::Any) = nothing
 
 Base.conj(x::Wirtinger) = error("`conj(::Wirtinger)` not yet defined")
 
-function add_wirtinger(a::Wirtinger, b::Wirtinger)
-    return Wirtinger(add(a.primal, b.primal), add(a.conjugate, b.conjugate))
-end
-
-add_wirtinger(a::Wirtinger, b) = add(a, Wirtinger(b, Zero()))
-add_wirtinger(a, b::Wirtinger) = add(Wirtinger(a, Zero()), b)
-
-function mul_wirtinger(a::Wirtinger, b::Wirtinger)
-    error("""
-          cannot multiply two Wirtinger objects; this error likely means a
-          `WirtingerRule` was inappropriately defined somewhere. Multiplication
-          of two Wirtinger objects is not defined because chain rule application
-          often expands into a non-commutative operation in the Wirtinger
-          calculus. To put it another way: simply given two Wirtinger objects
-          and no other information, we can't know "locally" which components to
-          conjugate in order to implement the chain rule. We could pick a
-          convention; for example, we could define `a::Wirtinger * b::Wirtinger`
-          such that we assume the chain rule application is of the form `f_a ∘ f_b`
-          instead of `f_b ∘ f_a`. However, picking such a convention is likely to
-          lead to silently incorrect derivatives due to commutativity assumptions
-          in downstream generic code that deals with the reals. Thus, ChainRulesCore
-          makes this operation an error instead.
-          """)
-end
-
-mul_wirtinger(a::Wirtinger, b) = Wirtinger(mul(a.primal, b), mul(a.conjugate, b))
-mul_wirtinger(a, b::Wirtinger) = Wirtinger(mul(a, b.primal), mul(a, b.conjugate))
 
 #####
 ##### `Casted`
@@ -174,14 +115,6 @@ Base.iterate(x::Casted, state) = iterate(x.value, state)
 
 Base.conj(x::Casted) = cast(conj, x.value)
 
-add_casted(a::Casted, b::Casted) = Casted(broadcasted(add, a.value, b.value))
-add_casted(a::Casted, b) = Casted(broadcasted(add, a.value, b))
-add_casted(a, b::Casted) = Casted(broadcasted(add, a, b.value))
-
-mul_casted(a::Casted, b::Casted) = Casted(broadcasted(mul, a.value, b.value))
-mul_casted(a::Casted, b) = Casted(broadcasted(mul, a.value, b))
-mul_casted(a, b::Casted) = Casted(broadcasted(mul, a, b.value))
-
 #####
 ##### `Zero`
 #####
@@ -200,13 +133,6 @@ Base.Broadcast.broadcastable(::Zero) = Ref(Zero())
 Base.iterate(x::Zero) = (x, nothing)
 Base.iterate(::Zero, ::Any) = nothing
 
-add_zero(::Zero, ::Zero) = Zero()
-add_zero(::Zero, b) = b
-add_zero(a, ::Zero) = a
-
-mul_zero(::Zero, ::Zero) = Zero()
-mul_zero(::Zero, ::Any) = Zero()
-mul_zero(::Any, ::Zero) = Zero()
 
 #####
 ##### `DNE`
@@ -228,14 +154,6 @@ Base.Broadcast.broadcastable(::DNE) = Ref(DNE())
 Base.iterate(x::DNE) = (x, nothing)
 Base.iterate(::DNE, ::Any) = nothing
 
-add_dne(::DNE, ::DNE) = DNE()
-add_dne(::DNE, b) = b
-add_dne(a, ::DNE) = a
-
-mul_dne(::DNE, ::DNE) = DNE()
-mul_dne(::DNE, ::Any) = DNE()
-mul_dne(::Any, ::DNE) = DNE()
-
 #####
 ##### `One`
 #####
@@ -254,13 +172,6 @@ Base.Broadcast.broadcastable(::One) = Ref(One())
 Base.iterate(x::One) = (x, nothing)
 Base.iterate(::One, ::Any) = nothing
 
-add_one(a::One, b::One) = add(extern(a), extern(b))
-add_one(a::One, b) = add(extern(a), b)
-add_one(a, b::One) = add(a, extern(b))
-
-mul_one(::One, ::One) = One()
-mul_one(::One, b) = b
-mul_one(a, ::One) = a
 
 #####
 ##### `Thunk`
@@ -295,11 +206,3 @@ end
 end
 
 Base.conj(x::Thunk) = @thunk(conj(extern(x)))
-
-add_thunk(a::Thunk, b::Thunk) = add(extern(a), extern(b))
-add_thunk(a::Thunk, b) = add(extern(a), b)
-add_thunk(a, b::Thunk) = add(a, extern(b))
-
-mul_thunk(a::Thunk, b::Thunk) = mul(extern(a), extern(b))
-mul_thunk(a::Thunk, b) = mul(extern(a), b)
-mul_thunk(a, b::Thunk) = mul(a, extern(b))

src/rule_definition_tools.jl

@@ -114,8 +114,8 @@ function rule_from_partials(input_arg, ∂s...)
     Δs = [Symbol(string(:Δ, i)) for i in 1:length(∂s)]
     Δs_tuple = Expr(:tuple, Δs...)
     if isempty(wirtinger_indices)
-        ∂_mul_Δs = [:(mul(@thunk($(∂s[i])), $(Δs[i]))) for i in 1:length(∂s)]
-        return :(Rule($Δs_tuple -> add($(∂_mul_Δs...))))
+        ∂_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[]
@@ -125,20 +125,20 @@ function rule_from_partials(input_arg, ∂s...)
                 Δi = Δs[i]
                 ∂i = Symbol(string(:∂, i))
                 push!(∂_wirtinger_defs, :($∂i = $(∂s[i])))
-                ∂f∂i_mul_Δ = :(mul(wirtinger_primal($∂i), wirtinger_primal($Δi)))
-                ∂f∂ī_mul_Δ̄ = :(mul(conj(wirtinger_conjugate($∂i)), wirtinger_conjugate($Δi)))
-                ∂f̄∂i_mul_Δ = :(mul(wirtinger_conjugate($∂i), wirtinger_primal($Δi)))
-                ∂f̄∂ī_mul_Δ̄ = :(mul(conj(wirtinger_primal($∂i)), wirtinger_conjugate($Δi)))
-                push!(∂_mul_Δs_primal, :(add($∂f∂i_mul_Δ, $∂f∂ī_mul_Δ̄)))
-                push!(∂_mul_Δs_conjugate, :(add($∂f̄∂i_mul_Δ, $∂f̄∂ī_mul_Δ̄)))
+                ∂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_Δ = :(mul(@thunk($(∂s[i])), $(Δs[i])))
+                ∂_mul_Δ = :(@thunk($(∂s[i])) * $(Δs[i]))
                 push!(∂_mul_Δs_primal, ∂_mul_Δ)
                 push!(∂_mul_Δs_conjugate, ∂_mul_Δ)
             end
         end
-        primal_rule = :(Rule($Δs_tuple -> add($(∂_mul_Δs_primal...))))
-        conjugate_rule = :(Rule($Δs_tuple -> add($(∂_mul_Δs_conjugate...))))
+        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)

src/rule_types.jl

@@ -78,7 +78,7 @@ accumulate(Δ, rule::Rule{typeof(df)}, x) = # customized `accumulate` implementa
 
 See also: [`accumulate!`](@ref), [`store!`](@ref), [`AbstractRule`](@ref)
 """
-accumulate(Δ, rule::AbstractRule, args...) = add(Δ, rule(args...))
+accumulate(Δ, rule::AbstractRule, args...) = Δ + rule(args...)
 
 """
     accumulate!(Δ, rule::AbstractRule, args...)
@@ -91,7 +91,7 @@ Note that this function internally calls `Base.Broadcast.materialize!(Δ, ...)`.
 See also: [`accumulate`](@ref), [`store!`](@ref), [`AbstractRule`](@ref)
 """
 function accumulate!(Δ, rule::AbstractRule, args...)
-    return materialize!(Δ, broadcastable(add(cast(Δ), rule(args...))))
+    return materialize!(Δ, broadcastable(cast(Δ) + rule(args...)))
 end
 
 accumulate!(Δ::Number, rule::AbstractRule, args...) = accumulate(Δ, rule, args...)
@@ -116,15 +116,6 @@ store!(Δ, rule::AbstractRule, args...) = materialize!(Δ, broadcastable(rule(ar
 ##### `Rule`
 #####
 
-Cassette.@context RuleContext
-
-const RULE_CONTEXT = Cassette.disablehooks(RuleContext())
-
-Cassette.overdub(::RuleContext, ::typeof(+), a, b) = add(a, b)
-Cassette.overdub(::RuleContext, ::typeof(*), a, b) = mul(a, b)
-
-Cassette.overdub(::RuleContext, ::typeof(add), a, b) = add(a, b)
-Cassette.overdub(::RuleContext, ::typeof(mul), a, b) = mul(a, b)
 
 """
     Rule(propation_function[, updating_function])
@@ -158,7 +149,7 @@ end
 # constructor based on a `UnionAll`, we get `Rule{Type{Thing}}` instead of `Rule{UnionAll}`
 Rule(f) = Rule{Core.Typeof(f),Nothing}(f, nothing)
 
-(rule::Rule{F})(args...) where {F} = Cassette.overdub(RULE_CONTEXT, rule.f, args...)
+(rule::Rule)(args...) = rule.f(args...)
 
 Base.show(io::IO, rule::Rule{<:Any, Nothing}) = print(io, "Rule($(rule.f))")
 Base.show(io::IO, rule::Rule) = print(io, "Rule($(rule.f), $(rule.u))")
@@ -212,7 +203,7 @@ otherwise return `WirtingerRule(P, C)`.
 """
 function AbstractRule(𝒟::Type, primal::AbstractRule, conjugate::AbstractRule)
     if 𝒟 <: Real || eltype(𝒟) <: Real
-        return Rule((args...) -> add(primal(args...), conjugate(args...)))
+        return Rule((args...) -> (primal(args...) + conjugate(args...)))
     else
         return WirtingerRule(primal, conjugate)
     end