Overhaul Rules (partner PR)

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "0.1.1"
+version = "0.2.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -10,7 +10,7 @@ Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
-ChainRulesCore = "^0.2"
+ChainRulesCore = "^0.3"
 FiniteDifferences = "^0.7"
 julia = "^1.0"
 

src/helper_functions.jl

@@ -1,8 +1,4 @@
-# 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.
+# Internal helpers for defining the `add!` field of an `InplaceableThunk`
 
 _update!(x, y) = x + y
 _update!(x::Array{T,N}, y::AbstractArray{T,N}) where {T,N} = x .+= y
@@ -11,20 +7,22 @@ _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),))
+    y = extern(y)
+    yp = getproperty(y, p)
+    xp = getproperty(x, p)
+    new_xp = _update!(xp, yp)
+    new = NamedTuple{(p,)}((new_xp,))
     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
-
+"""
+    _checked_rrule
 
+like `rrule` but throws an error if the `rrule` is not defined.
+Rather than returning `nothing`
+"""
 function _checked_rrule(f, args...; kwargs...)
     r = rrule(f, args...; kwargs...)
     r isa Nothing && _throw_checked_rrule_error(f, args...; kwargs...)

src/rulesets/Base/array.jl

@@ -3,56 +3,74 @@
 #####
 
 function rrule(::typeof(reshape), A::AbstractArray, dims::Tuple{Vararg{Int}})
-    return reshape(A, dims), (Rule(Ȳ->reshape(Ȳ, dims)), DNERule())
+    function reshape_pullback(Ȳ)
+        return (NO_FIELDS, @thunk(reshape(Ȳ, dims)), DNE())
+    end
+    return reshape(A, dims), reshape_pullback
 end
 
 function rrule(::typeof(reshape), A::AbstractArray, dims::Int...)
-    Y, (rule, _) = rrule(reshape, A, dims)
-    return Y, (rule, fill(DNERule(), length(dims))...)
+    function reshape_pullback(Ȳ)
+        ∂A = @thunk(reshape(Ȳ, dims))
+        return (NO_FIELDS, ∂A, fill(DNE(), length(dims))...)
+    end
+    return reshape(A, dims...), reshape_pullback
 end
 
 #####
 ##### `hcat` (🐈)
 #####
 
 function rrule(::typeof(hcat), A::AbstractArray, Bs::AbstractArray...)
-    Y = hcat(A, Bs...)
-    Xs = (A, Bs...)
-    rules = ntuple(length(Bs) + 1) do i
-        l = mapreduce(j->size(Xs[j], 2), Base.add_sum, 1:i-1; init=0)
-        u = l + size(Xs[i], 2)
-        dim = u > l + 1 ? (l+1:u) : u
-        # NOTE: The copy here is defensive, since `selectdim` returns a view which we can
-        # materialize with `copy`
-        Rule(Ȳ->copy(selectdim(Ȳ, 2, dim)))
+    function hcat_pullback(Ȳ)
+        Xs = (A, Bs...)
+        ntuple(length(Bs) + 2) do full_i
+            full_i == 1 && return NO_FIELDS
+
+            i = full_i - 1
+            l = mapreduce(j->size(Xs[j], 2), Base.add_sum, 1:i-1; init=0)
+            u = l + size(Xs[i], 2)
+            dim = u > l + 1 ? (l+1:u) : u
+            # NOTE: The copy here is defensive, since `selectdim` returns a view which we can
+            # materialize with `copy`
+            copy(selectdim(Ȳ, 2, dim))
+        end
     end
-    return Y, rules
+    return hcat(A, Bs...), hcat_pullback
 end
 
 #####
 ##### `vcat`
 #####
 
 function rrule(::typeof(vcat), A::AbstractArray, Bs::AbstractArray...)
-    Y = vcat(A, Bs...)
-    n = size(A, 1)
-    ∂A = Rule(Ȳ->copy(selectdim(Ȳ, 1, 1:n)))
-    ∂Bs = ntuple(length(Bs)) do i
-        l = n + mapreduce(j->size(Bs[j], 1), Base.add_sum, 1:i-1; init=0)
-        u = l + size(Bs[i], 1)
-        Rule(Ȳ->copy(selectdim(Ȳ, 1, l+1:u)))
+    function vcat_pullback(Ȳ)
+        n = size(A, 1)
+        ∂A = copy(selectdim(Ȳ, 1, 1:n))
+        ∂Bs = ntuple(length(Bs)) do i
+            l = n + mapreduce(j->size(Bs[j], 1), Base.add_sum, 1:i-1; init=0)
+            u = l + size(Bs[i], 1)
+            copy(selectdim(Ȳ, 1, l+1:u))
+        end
+        return (NO_FIELDS, ∂A, ∂Bs...)
     end
-    return Y, (∂A, ∂Bs...)
+    return vcat(A, Bs...), vcat_pullback
 end
 
 #####
 ##### `fill`
 #####
 
 function rrule(::typeof(fill), value::Any, dims::Tuple{Vararg{Int}})
-    return fill(value, dims), (Rule(sum), DNERule())
+    function fill_pullback(Ȳ)
+        return (NO_FIELDS, @thunk(sum(Ȳ)), DNE())
+    end
+    return fill(value, dims), fill_pullback
 end
 
 function rrule(::typeof(fill), value::Any, dims::Int...)
-    return fill(value, dims), (Rule(sum), ntuple(_->DNERule(), length(dims))...)
+    function fill_pullback(Ȳ)
+        return (NO_FIELDS, @thunk(sum(Ȳ)), ntuple(_->DNE(), length(dims))...)
+    end
+    return fill(value, dims), fill_pullback
 end

src/rulesets/Base/base.jl

@@ -103,10 +103,30 @@
 
 # product rule requires special care for arguments where `mul` is non-commutative
 
-frule(::typeof(*), x::Number, y::Number) = x * y, Rule((Δx, Δy) -> Δx * y + x * Δy)
-
-rrule(::typeof(*), x::Number, y::Number) = x * y, (Rule(ΔΩ -> ΔΩ * y'), Rule(ΔΩ -> x' * ΔΩ))
-
-frule(::typeof(identity), x) = x, Rule(identity)
-
-rrule(::typeof(identity), x) = x, Rule(identity)
+function frule(::typeof(*), x::Number, y::Number)
+    function times_pushforward(_, Δx, Δy)
+        return Δx * y + x * Δy
+    end
+    return x * y, times_pushforward
+end
+
+function rrule(::typeof(*), x::Number, y::Number)
+    function times_pullback(ΔΩ)
+        return (NO_FIELDS,  @thunk(ΔΩ * y'), @thunk(x' * ΔΩ))
+    end
+    return x * y, times_pullback
+end
+
+function frule(::typeof(identity), x)
+    function identity_pushforward(_, ẏ)
+        return ẏ
+    end
+    return x, identity_pushforward
+end
+
+function rrule(::typeof(identity), x)
+    function identity_pullback(ȳ)
+        return (NO_FIELDS, ȳ)
+    end
+    return x, identity_pullback
+end

src/rulesets/Base/broadcast.jl

@@ -5,20 +5,26 @@ without relying on inference hacks unless we have something akin to
 https://github.com/JuliaLang/julia/issues/22129.
 =#
 function _cast_diff(f, x)
-    element_rule = u -> begin
+    function element_rule(u)
         fu, du = frule(f, u)
-        fu, extern(du(One()))
+        fu, extern(du(NamedTuple(), One()))
     end
     results = broadcast(element_rule, x)
     return first.(results), last.(results)
 end
 
 function frule(::typeof(broadcast), f, x)
     Ω, ∂x = _cast_diff(f, x)
-    return Ω, Rule((_, Δx) -> Δx * cast(∂x))
+    function broadcast_pushforward(_, Δf, Δx)
+        return Δx * cast(∂x)
+    end
+    return Ω, broadcast_pushforward
 end
 
 function rrule(::typeof(broadcast), f, x)
     values, derivs = _cast_diff(f, x)
-    return values, (DNERule(), Rule(ΔΩ -> ΔΩ * cast(derivs)))
+    function broadcast_pullback(ΔΩ)
+        return (NO_FIELDS, DNE(), @thunk(ΔΩ * cast(derivs)))
+    end
+    return values, broadcast_pullback
 end

src/rulesets/Base/mapreduce.jl

@@ -4,15 +4,19 @@
 
 function rrule(::typeof(map), f, xs...)
     y = map(f, xs...)
-    ∂xs = ntuple(length(xs)) do i
-        Rule() do ȳ
-            map(ȳ, xs...) do ȳi, xis...
-                _, ∂xis = _checked_rrule(f, xis...)
-                extern(∂xis[i](ȳi))
+    function map_pullback(ȳ)
+        ntuple(length(xs)+2) do full_i
+            full_i == 1 && return NO_FIELDS
+            full_i == 2 && return DNE()
+            i = full_i-2
+            @thunk map(ȳ, xs...) do ȳi, xis...
+                _, pullback = _checked_rrule(f, xis...)
+                ∂xis = pullback(ȳi)
+                extern(∂xis[i+1])  #+1 to skp ∂self
             end
         end
     end
-    return y, (DNERule(), ∂xs...)
+    return y, map_pullback
 end
 
 #####
@@ -26,15 +30,18 @@ for mf in (:mapreduce, :mapfoldl, :mapfoldr)
         insert!(sig.args, 2, Expr(:parameters, Expr(:kw, :dims, :(:))))
         insert!(call.args, 2, Expr(:parameters, Expr(:kw, :dims, :dims)))
     end
+    pullback_name = Symbol(mf, :_pullback)
     body = quote
         y = $call
-        ∂x = Rule() do ȳ
-            broadcast(x, ȳ) do xi, ȳi
-                _, ∂xi = _checked_rrule(f, xi)
-                extern(∂xi(ȳi))
+        function $pullback_name(ȳ)
+            ∂x = @thunk broadcast(x, ȳ) do xi, ȳi
+                _, pullback_f = _checked_rrule(f, xi)
+                _, ∂xi = pullback_f(ȳi)
+                extern(∂xi)
             end
+            (NO_FIELDS, DNE(), DNE(), ∂x)
         end
-        return y, (DNERule(), DNERule(), ∂x)
+        return y, $pullback_name
     end
     eval(Expr(:function, sig, body))
 end
@@ -43,22 +50,40 @@ end
 ##### `sum`
 #####
 
-frule(::typeof(sum), x) = (sum(x), Rule(sum))
+function frule(::typeof(sum), x)
+    function sum_pushforward(_, ẋ)
+        return sum(ẋ)
+    end
+    return sum(x), sum_pushforward
+end
 
-rrule(::typeof(sum), x) = (sum(x), Rule(cast))
+function rrule(::typeof(sum), x)
+    function sum_pullback(ȳ)
+        return (NO_FIELDS, cast(ȳ))
+    end
+    return sum(x), sum_pullback
+end
 
 function rrule(::typeof(sum), f, x::AbstractArray{<:Real}; dims=:)
-    y, (_, _, ∂x) = rrule(mapreduce, f, Base.add_sum, x; dims=dims)
-    return y, (DNERule(), ∂x)
+    y, mr_pullback = rrule(mapreduce, f, Base.add_sum, x; dims=dims)
+    function sum_pullback(ȳ)
+        NO_FIELDS, DNE(), last(mr_pullback(ȳ))
+    end
+    return y, sum_pullback
 end
 
 function rrule(::typeof(sum), x::AbstractArray{<:Real}; dims=:)
-    y, (_, ∂x) = rrule(sum, identity, x; dims=dims)
-    return y, ∂x
+    y,  inner_pullback = rrule(sum, identity, x; dims=dims)
+    function sum_pullback(ȳ)
+        NO_FIELDS, last(inner_pullback(ȳ))
+    end
+    return y, sum_pullback
 end
 
 function rrule(::typeof(sum), ::typeof(abs2), x::AbstractArray{<:Real}; dims=:)
     y = sum(abs2, x; dims=dims)
-    ∂x = Rule(ȳ -> 2ȳ .* x)
-    return y, (DNERule(), ∂x)
+    function sum_abs2_pullback(ȳ)
+        return (NO_FIELDS, DNE(), @thunk(2ȳ .* x))
+    end
+    return y, sum_abs2_pullback
 end