Merge pull request #180 from KDr2/chainrules

Integrate ReverseDiff with ChainRules

Author: yebai

Project.toml

@@ -1,8 +1,9 @@
 name = "ReverseDiff"
 uuid = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-version = "1.9.0"
+version = "1.10.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -16,6 +17,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+ChainRulesCore = "1"
 DiffResults = "1"
 DiffRules = "0.1, 1"
 ForwardDiff = "0.10"

src/ReverseDiff.jl

@@ -17,6 +17,8 @@ using StaticArrays
 
 using MacroTools
 
+using ChainRulesCore
+
 # Not all operations will be valid over all of these types, but that's okay; such cases
 # will simply error when they hit the original operation in the overloaded definition.
 const ARRAY_TYPES = (:AbstractArray, :AbstractVector, :AbstractMatrix, :Array, :Vector, :Matrix)

src/macros.jl

@@ -237,7 +237,136 @@ macro grad(expr)
         end
     end |> esc
 end
+
+"""
+    _make_fwd_args(func, arg_list)
+
+Function `_make_fwd_args` accepts a function name and an argument
+list, returns a tuple of argument lists whose elements are:
+1. the`arg_list` untouched, 2. a new argument list with the function
+as its first element and other elements in `arg_list` followed, 3. a
+new argument for the definition of function `track`, 4. a new argument
+list with all kwargs removed, 5, types of the arguments in the 4th
+element, 5 the kwargs name if any otherwise an empty tuple. E.g.:
+
+_make_fwd_args(:f, [:(a::String), :(b::TrackedReal), :(args...)])
+
+returns
+
+([:(a::String), :(b::TrackedReal), :(args...)],
+ [:f, :(a::String), :(b::TrackedReal), :(args...)],
+ [:(::typeof(f)), :(a::String), :(b::TrackedReal), :(args...)],
+ [:(a::String), :(b::TrackedReal), :(args...)],
+ [:String, :TrackedReal, :(Vararg{Any})],
+ :kwargs)
+
+It also deals with varargs and variable keyword arguments, and ensures
+that at least one of the argument is tracked.
+
+"""
+function _make_fwd_args(func, args_l)
+    kwargs = :(())
+    args_r = copy(args_l)
+    args_track = copy(args_l)
+    if Meta.isexpr(args_r[1], :parameters) # has kw args
+        insert!(args_r, 2, func)
+        insert!(args_track, 2, :(::typeof($func)))
+        kwargs = gensym(:kwargs)
+        args_track[1].args = [:($(kwargs)...)]
+    else
+        insert!(args_r, 1, func)
+        insert!(args_track, 1, :(::typeof($func)))
+    end
+
+    args_fixed = filter(copy(args_l)) do arg
+        !Meta.isexpr(arg, :parameters)
+    end
+
+    arg_types = map(args_fixed) do arg
+        if Meta.isexpr(arg, :(...))
+            Meta.isexpr(arg.args[1], :(::)) ? :(Vararg{$(arg.args[1].args[end])}) : :(Vararg{Any})
+        elseif Meta.isexpr(arg, :(::))
+            arg.args[end]
+        else
+            :Any
+        end
+    end
+
+    return args_l, args_r, args_track, args_fixed, arg_types, kwargs
+end
+
+"""
+    @grad_from_chainrules f(args...; kwargs...)
+
+The `@grad_from_chainrules` macro provides a way to import
+adjoints(rrule) defined in ChainRules to ReverseDiff. One must provide
+a method signature to import the corresponding `rrule`. In the
+provided method signature, one should replace the types of arguments
+to which one wants to take derivatives with respect with
+`ReverseDiff.TrackedReal` and `ReverseDiff.TrackedArray`
+respectively. For example, we can import `rrule` of `f(x::Real,
+y::Array)` like below:
+
+```julia
+ReverseDiff.@grad_from_chainrules f(x::TrackedReal, y::TrackedArray)
+ReverseDiff.@grad_from_chainrules f(x::TrackedReal, y::Array)
+ReverseDiff.@grad_from_chainrules f(x::Real, y::TrackedArray)
+```
+"""
+macro grad_from_chainrules(fcall)
+    Meta.isexpr(fcall, :call) && length(fcall.args) >= 2 ||
+        error("`@grad_from_chainrules` has to be applied to a function signature")
+    f = esc(fcall.args[1])
+    xs = fcall.args[2:end]
+    args_l, args_r, args_track, args_fixed, arg_types, kwargs = _make_fwd_args(f, xs)
+
+    return quote
+        $f($(args_l...)) = ReverseDiff.track($(args_r...))
+        function ReverseDiff.track($(args_track...))
+            args = ($(args_fixed...),)
+            tp = ReverseDiff.tape(args...)
+            output_value, back = ChainRulesCore.rrule($f, map(ReverseDiff.value, args)...; $kwargs...)
+            output = ReverseDiff.track(output_value, tp)
+            closure(cls_args...; cls_kwargs...) = ChainRulesCore.rrule($f, map(ReverseDiff.value, cls_args)...; cls_kwargs...)
+            ReverseDiff.record!(
+                tp,
+                ReverseDiff.SpecialInstruction,
+                $f,
+                args,
+                output,
+                (back, closure, $kwargs),
+            )
+            return output
+        end
+
+        @noinline function ReverseDiff.special_reverse_exec!(instruction::ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}})
+            output = instruction.output
+            input = instruction.input
+            back = instruction.cache[1]
+            back_output = back(ReverseDiff.deriv(output))
+            input_derivs = back_output[2:end]
+            @assert input_derivs isa Tuple
+            ReverseDiff._add_to_deriv!.(input, input_derivs)
+            ReverseDiff.unseed!(output)
+            return nothing
+        end
+
+        @noinline function ReverseDiff.special_forward_exec!(instruction::ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}})
+            output, input = instruction.output, instruction.input
+            ReverseDiff.pull_value!.(input)
+            pullback = instruction.cache[2]
+            kwargs = instruction.cache[3]
+            out_value = pullback(input...; kwargs...)[1]
+            ReverseDiff.value!(output, out_value)
+            return nothing
+        end
+    end
+end
+
 _add_to_deriv!(d1, d2) = nothing
+function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray{<:TrackedReal}}, d2::AbstractThunk)
+    increment_deriv!(d1, unthunk(d2))
+end
 function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray{<:TrackedReal}}, d2)
     increment_deriv!(d1, d2)
 end

test/ChainRulesTests.jl

@@ -0,0 +1,210 @@
+module ChainRulesTest
+
+using LinearAlgebra
+using ChainRulesCore
+using DiffResults
+using ReverseDiff
+using Test
+
+f(x) = sum(4x .+ 1)
+
+function ChainRulesCore.rrule(::typeof(f), x)
+    r = f(x)
+    function back(d)
+        #=
+        The proper derivative of `f` is 4, but in order to
+        check if `ChainRulesCore.rrule` had taken over the compuation,
+        we define a rrule that returns 3 as `f`'s derivative.
+
+        After importing this rrule into ReverseDiff, if we get 3
+        rather than 4 when we compute the derivative of `f`, it means
+        the importing mechanism works.
+        =#
+        return ChainRulesCore.NoTangent(), fill(3 * d, size(x))
+    end
+    return r, back
+end
+
+ReverseDiff.@grad_from_chainrules f(x::ReverseDiff.TrackedArray)
+
+
+g(x, y) = sum(4x .+ 4y)
+
+function ChainRulesCore.rrule(::typeof(g), x, y)
+    r = g(x, y)
+    function back(d)
+        # same as above, use 3 and 5 as the derivatives
+        return ChainRulesCore.NoTangent(), fill(3 * d, size(x)), fill(5 * d, size(x))
+    end
+    return r, back
+end
+
+ReverseDiff.@grad_from_chainrules g(x::ReverseDiff.TrackedArray, y)
+ReverseDiff.@grad_from_chainrules g(x, y::ReverseDiff.TrackedArray)
+ReverseDiff.@grad_from_chainrules g(x::ReverseDiff.TrackedArray, y::ReverseDiff.TrackedArray)
+
+@testset "rrule in ChainRules and ReverseDiff" begin
+    ## ChainRules
+    # function f
+    input = rand(3, 3)
+    output, back = ChainRulesCore.rrule(f, input);
+    _, d = back(1)
+    @test output == f(input)
+    @test d == fill(3, size(input))
+    # function g
+    inputs = rand(3, 3), rand(3, 3)
+    output, back = ChainRulesCore.rrule(g, inputs...);
+    _, d1, d2 = back(1)
+    @test output == g(inputs...)
+    @test d1 == fill(3, size(inputs[1]))
+    @test d2 == fill(5, size(inputs[2]))
+
+
+    ## ReverseDiff
+    #function f
+    inputs = (rand(3, 3), )
+
+    results = (similar(inputs[1]),)
+    f_tape = ReverseDiff.GradientTape(x -> f(x) + 2, (rand(3, 3),))
+    ReverseDiff.gradient!(results, f_tape, inputs)
+
+    @test results[1] == fill(3, size(inputs[1]))
+
+    results = (similar(inputs[1]),)
+    compiled_tape = ReverseDiff.CompiledTape(f_tape)
+    ReverseDiff.gradient!(results, compiled_tape, inputs)
+    @test results[1] == fill(3, size(inputs[1]))
+
+    # function g
+    inputs = rand(3, 3), rand(3, 3)
+
+    results = (similar(inputs[1]), similar(inputs[2]))
+    f_tape = ReverseDiff.GradientTape((x, y) -> g(x, y) + 2, (rand(3, 3), rand(3, 3)))
+    ReverseDiff.gradient!(results, f_tape, inputs)
+
+    @test results[1] == fill(3, size(inputs[1]))
+    @test results[2] == fill(5, size(inputs[2]))
+
+    results = (similar(inputs[1]), similar(inputs[2]),)
+    compiled_tape = ReverseDiff.CompiledTape(f_tape)
+    ReverseDiff.gradient!(results, compiled_tape, inputs)
+    @test results[1] == fill(3, size(inputs[1]))
+    @test results[2] == fill(5, size(inputs[2]))
+
+end
+
+### Tape test
+@testset "Tape test: Ensure ordinary call is not tracked" begin
+    tp = ReverseDiff.InstructionTape()
+
+    f(x) = sum(2x .+ g([1, 2], [3, 4]))
+    x = rand(3, 3)
+    xt = ReverseDiff.track(copy(x), tp)
+    # record
+    yt = f(xt)
+    @test length(tp) == 3 # sum, broadcast+, broadcast*, but not `g`
+end
+
+### Functions with varargs and kwargs
+# Varargs
+f_vararg(x, args...) = sum(4x .+ sum(args))
+
+function ChainRulesCore.rrule(::typeof(f_vararg), x, args...)
+    r = f_vararg(x, args...)
+    function back(d)
+        return ChainRulesCore.NoTangent(), fill(3 * d, size(x))
+    end
+    return r, back
+end
+
+ReverseDiff.@grad_from_chainrules f_vararg(x::ReverseDiff.TrackedArray, args...)
+
+@testset "Function with Varargs" begin
+    inputs = (rand(3, 3), )
+
+    results = (similar(inputs[1]),)
+    f_tape = ReverseDiff.GradientTape(x -> f_vararg(x, 1, 2, 3) + 2, (rand(3, 3),))
+    ReverseDiff.gradient!(results, f_tape, inputs)
+
+    @test results[1] == fill(3, size(inputs[1]))
+end
+
+
+# Vargs and kwargs
+f_kw(x, args...; k=1, kwargs...) = sum(4x .+ sum(args) .+ (k + kwargs[:j]))
+
+function ChainRulesCore.rrule(::typeof(f_kw), x, args...; k=1, kwargs...)
+    r = f_kw(x, args...; k=k, kwargs...)
+    function back(d)
+        return ChainRulesCore.NoTangent(), fill(3 * d, size(x))
+    end
+    return r, back
+end
+
+ReverseDiff.@grad_from_chainrules f_kw(x::ReverseDiff.TrackedArray, args...; k=1, kwargs...)
+
+@testset "Function with Varargs and kwargs" begin
+    inputs = (rand(3, 3), )
+
+    results = (similar(inputs[1]),)
+    f_tape = ReverseDiff.GradientTape(x -> f_kw(x, 1, 2, 3; k=2, j=3) + 2, (rand(3, 3),))
+    ReverseDiff.gradient!(results, f_tape, inputs)
+
+    @test results[1] == fill(3, size(inputs[1]))
+end
+
+### Mix @grad and @grad_from_chainrules
+
+h(x) = 10x
+h(x::ReverseDiff.TrackedArray) = ReverseDiff.track(h, x)
+ReverseDiff.@grad function h(x)
+    xv = ReverseDiff.value(x)
+    return h(xv), Δ -> (Δ * 7,) # use 7 asits derivatives
+end
+
+@testset "ReverseDiff and ChainRules Mixed" begin
+    t(x) = g(x, h(x))
+    inputs = (rand(3, 3), )
+    results = (similar(inputs[1]),)
+
+    g_tape = ReverseDiff.GradientTape(t, (rand(3, 3),))
+    ReverseDiff.gradient!(results, g_tape, inputs)
+    @test results[1] == fill(38, size(inputs[1])) # 38 = 3 + 5 * 7
+end
+
+### Isolated Scope
+module IsolatedModuleForTestingScoping
+using ChainRulesCore
+using ReverseDiff: @grad_from_chainrules
+
+f(x) = sum(4x .+ 1)
+
+function ChainRulesCore.rrule(::typeof(f), x)
+    r = f(x)
+    function back(d)
+        # return a distinguishable but improper grad
+        return ChainRulesCore.NoTangent(), fill(3 * d, size(x))
+    end
+    return r, back
+end
+
+@grad_from_chainrules f(x::TrackedArray)
+
+module SubModule
+using Test
+using ReverseDiff: TrackedArray, GradientTape, gradient!
+using ..IsolatedModuleForTestingScoping: f
+@testset "rrule in Isolated Scope" begin
+    inputs = (rand(3, 3), )
+
+    results = (similar(inputs[1]),)
+    f_tape = GradientTape(x -> f(x) + 2, (rand(3, 3),))
+    gradient!(results, f_tape, inputs)
+
+    @test results[1] == fill(3, size(inputs[1]))
+end
+
+end # end of SubModule
+end # end of IsolatedModuleForTestingScoping
+
+end

test/derivatives/LinAlgTests.jl

@@ -79,7 +79,7 @@ function test_arr2arr(f, a, b, tp)
     ReverseDiff.value!(at, a2)
     ReverseDiff.forward_pass!(tp)
     @test value(ct) == f(a2, b)
-    
+
     ReverseDiff.value!(at, a)
     empty!(tp)
 

test/runtests.jl

@@ -14,6 +14,10 @@ println("running MacrosTests...")
 t = @elapsed include(joinpath(TESTDIR, "MacrosTests.jl"))
 println("done (took $t seconds).")
 
+println("running ChainRulesTests...")
+t = @elapsed include(joinpath(TESTDIR, "ChainRulesTests.jl"))
+println("done (took $t seconds).")
+
 println("running ScalarTests...")
 t = @elapsed include(joinpath(TESTDIR, "derivatives/ScalarTests.jl"))
 println("done (took $t seconds).")