make work for new ChainRules v0.1

Author: oxinabox

src/AbstractChainRules.jl

@@ -1,9 +1,10 @@
 module AbstractChainRules
-
 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("rules.jl")

src/rules.jl

@@ -112,32 +112,6 @@ See also: [`accumulate`](@ref), [`accumulate!`](@ref), [`AbstractRule`](@ref)
 """
 store!(Δ, rule::AbstractRule, args...) = materialize!(Δ, broadcastable(rule(args...)))
 
-# Special purpose updating for operations which can be done in-place. This function is
-# just internal and free-form; it is not a method of `accumulate!` directly as it does
-# not adhere to the expected method signature form, i.e. `accumulate!(value, rule, args)`.
-# Instead it's `_update!(old, new, extrastuff...)` and is not specific to any particular
-# rule.
-
-_update!(x, y) = x + y
-_update!(x::Array{T,N}, y::AbstractArray{T,N}) where {T,N} = x .+= y
-
-_update!(x, ::Zero) = x
-_update!(::Zero, y) = y
-_update!(::Zero, ::Zero) = Zero()
-
-function _update!(x::NamedTuple{Ns}, y::NamedTuple{Ns}) where Ns
-    return NamedTuple{Ns}(map(p->_update!(getproperty(x, p), getproperty(y, p)), Ns))
-end
-
-function _update!(x::NamedTuple, y, p::Symbol)
-    new = NamedTuple{(p,)}((_update!(getproperty(x, p), y),))
-    return merge(x, new)
-end
-
-function _update!(x::NamedTuple{Ns}, y::NamedTuple{Ns}, p::Symbol) where Ns
-    return _update!(x, getproperty(y, p), p)
-end
-
 #####
 ##### `Rule`
 #####
@@ -377,23 +351,6 @@ See also: [`frule`](@ref), [`AbstractRule`](@ref), [`@scalar_rule`](@ref)
 """
 rrule(::Any, ::Vararg{Any}; kwargs...) = nothing
 
-@noinline function _throw_checked_rrule_error(f, args...; kwargs...)
-    io = IOBuffer()
-    print(io, "can't differentiate `", f, '(')
-    join(io, map(arg->string("::", typeof(arg)), args), ", ")
-    if !isempty(kwargs)
-        print(io, ";")
-        join(io, map(((k, v),)->string(k, "=", v), kwargs), ", ")
-    end
-    print(io, ")`; no matching `rrule` is defined")
-    throw(ArgumentError(String(take!(io))))
-end
-
-function _checked_rrule(f, args...; kwargs...)
-    r = rrule(f, args...; kwargs...)
-    r isa Nothing && _throw_checked_rrule_error(f, args...; kwargs...)
-    return r
-end
 
 #####
 ##### macros

test/rules.jl

@@ -42,36 +42,4 @@ dummy_identity(x) = x
         @test rule[1] == rule
         @test_throws BoundsError rule[2]
     end
-    @testset "helper functions" begin
-        # Hits fallback, since we can't update `Diagonal`s in place
-        X = Diagonal([1, 1])
-        Y = copy(X)
-        @test AbstractChainRules._update!(X, [1 2; 3 4]) == [2 2; 3 5]
-        @test X == Y  # no change to X
-
-        X = [1 2; 3 4]
-        Y = copy(X)
-        @test AbstractChainRules._update!(X, Diagonal([1, 1])) == [2 2; 3 5]
-        @test X != Y  # X has been updated
-
-        # Reusing above X
-        @test AbstractChainRules._update!(X, Zero()) === X
-        @test AbstractChainRules._update!(Zero(), X) === X
-        @test AbstractChainRules._update!(Zero(), Zero()) === Zero()
-
-        X = (A=[1 0; 0 1], B=[2 2; 2 2])
-        Y = deepcopy(X)
-        @test AbstractChainRules._update!(X, Y) == (A=[2 0; 0 2], B=[4 4; 4 4])
-        @test X.A != Y.A
-        @test X.B != Y.B
-
-        try
-            # We defined a 2-arg method for `cool` but no `rrule`
-            AbstractChainRules._checked_rrule(cool, 1.0, 2.0)
-        catch e
-            @test e isa ArgumentError
-            @test e.msg == "can't differentiate `cool(::Float64, ::Float64)`; no " *
-                           "matching `rrule` is defined"
-        end
-    end
 end