Implement backend for ReverseDiff

Project.toml

@@ -14,14 +14,16 @@ Compat = "3"
 ExprTools = "0.1"
 ForwardDiff = "0.10"
 Requires = "0.5, 1"
+ReverseDiff = "1"
 julia = "1"
 
 [extras]
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "FiniteDifferences", "ForwardDiff", "Random", "Zygote"]
+test = ["Test", "FiniteDifferences", "ForwardDiff", "Random", "ReverseDiff", "Zygote"]

src/AbstractDifferentiation.jl

@@ -640,8 +640,12 @@ function zero_matrix_like(x)
     throw("The function `zero_matrix_like` is not defined for the type $(typeof(x)).")
 end
 
+@inline asarray(x) = [x]
+@inline asarray(x::AbstractArray) = x
+
 function __init__()
     @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include("forwarddiff.jl")
+    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" include("reversediff.jl")
     @require FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" include("finitedifferences.jl")
 end
 

src/forwarddiff.jl

@@ -76,9 +76,6 @@ end
 # support arrays and tuples
 @noinline step_toward(x, v, h) = x .+ h .* v
 
-@inline asarray(x) = [x]
-@inline asarray(x::AbstractArray) = x
-
 getchunksize(::Nothing) = Nothing
 getchunksize(::Val{N}) where {N} = N
 

src/reversediff.jl

@@ -0,0 +1,58 @@
+using .ReverseDiff: ReverseDiff, DiffResults
+
+primal_value(x::ReverseDiff.TrackedReal) = ReverseDiff.value(x)
+
+"""
+    ReverseDiffBackend
+
+AD backend that uses reverse mode with ReverseDiff.jl.
+"""
+struct ReverseDiffBackend <: AbstractReverseMode end
+
+@primitive function jacobian(ba::ReverseDiffBackend, f, xs...)
+    xs_arr = map(asarray, xs)
+    tape = ReverseDiff.JacobianTape(xs_arr) do (xs_arr...)
+        xs_new = map(xs, xs_arr) do x, x_arr
+            return x isa Number ? only(x_arr) : x_arr
+        end
+        return asarray(f(xs_new...))
+    end
+    results = ReverseDiff.jacobian!(tape, xs_arr)
+    return map(xs, results) do x, result
+        return x isa Number ? vec(result) : result
+    end
+end
+function jacobian(ba::ReverseDiffBackend, f, xs::AbstractArray...)
+    return ReverseDiff.jacobian(asarray ∘ f, xs)
+end
+
+function derivative(ba::ReverseDiffBackend, f, xs::Number...)
+    tape = ReverseDiff.InstructionTape()
+    xs_tracked = ReverseDiff.TrackedReal.(xs, zero.(xs), Ref(tape))
+    y_tracked = f(xs_tracked...)
+    ReverseDiff.seed!(y_tracked)
+    ReverseDiff.reverse_pass!(tape)
+    return ReverseDiff.deriv.(xs_tracked)
+end
+
+function gradient(ba::ReverseDiffBackend, f, xs::AbstractArray...)
+    return ReverseDiff.gradient(f, xs)
+end
+
+function hessian(ba::ReverseDiffBackend, f, x::AbstractArray)
+    return (ReverseDiff.hessian(f, x),)
+end
+
+function value_and_gradient(ba::ReverseDiffBackend, f, x::AbstractArray)
+    result = DiffResults.GradientResult(x)
+    cfg = ReverseDiff.GradientConfig(x)
+    ReverseDiff.gradient!(result, f, x, cfg)
+    return DiffResults.value(result), (DiffResults.derivative(result),)
+end
+
+function value_and_hessian(ba::ReverseDiffBackend, f, x)
+    result = DiffResults.HessianResult(x)
+    cfg = ReverseDiff.HessianConfig(result, x)
+    ReverseDiff.hessian!(result, f, x, cfg)
+    return DiffResults.value(result), (DiffResults.hessian(result),)
+end

test/reversediff.jl

@@ -0,0 +1,39 @@
+using AbstractDifferentiation
+using Test
+using ReverseDiff
+
+@testset "ReverseDiffBackend" begin
+    backends = [@inferred(AD.ReverseDiffBackend())]
+    @testset for backend in backends
+        @testset "Derivative" begin
+            test_derivatives(backend)
+        end
+        @testset "Gradient" begin
+            test_gradients(backend)
+        end
+        @testset "Jacobian" begin
+            test_jacobians(backend)
+        end
+        @testset "Hessian" begin
+            test_hessians(backend)
+        end
+        @testset "jvp" begin
+            test_jvp(backend)
+        end
+        @testset "j′vp" begin
+            test_j′vp(backend)
+        end
+        @testset "Lazy Derivative" begin
+            test_lazy_derivatives(backend)
+        end
+        @testset "Lazy Gradient" begin
+            test_lazy_gradients(backend)
+        end
+        @testset "Lazy Jacobian" begin
+            test_lazy_jacobians(backend)
+        end
+        @testset "Lazy Hessian" begin
+            test_lazy_hessians(backend)
+        end
+    end
+end

test/runtests.jl

@@ -5,5 +5,6 @@ using Test
     include("test_utils.jl")
     include("defaults.jl")
     include("forwarddiff.jl")
+    include("reversediff.jl")
     include("finitedifferences.jl")
 end