Add ReverseDiff support. (#23)

Add ReverseDiff support,fix #22, improve missing AD method message

Author: tpapp

Project.toml

@@ -21,8 +21,9 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Distributions", "Documenter", "Flux", "ForwardDiff", "Test", "StatsBase"]
+test = ["Distributions", "Documenter", "Flux", "ForwardDiff", "ReverseDiff", "Test", "StatsBase"]

src/AD_ForwardDiff.jl

@@ -13,12 +13,10 @@ function show(io::IO, ℓ::ForwardDiffLogDensity)
           ", w/ chunk size ", length(ℓ.gradientconfig.seeds))
 end
 
-_anyargument(ℓ) = zeros(dimension(ℓ))
-
 _default_chunk(ℓ) = ForwardDiff.Chunk(dimension(ℓ))
 
 _default_gradientconfig(ℓ, chunk::ForwardDiff.Chunk) =
-    ForwardDiff.GradientConfig(_value_closure(ℓ), _anyargument(ℓ), chunk)
+    ForwardDiff.GradientConfig(_value_closure(ℓ), _vectorargument(ℓ), chunk)
 
 _default_gradientconfig(ℓ, chunk::Integer) =
     _default_gradientconfig(ℓ, ForwardDiff.Chunk(chunk))
@@ -42,7 +40,7 @@ end
 
 function logdensity(::Type{ValueGradient}, fℓ::ForwardDiffLogDensity, x::RealVector)
     @unpack ℓ, gradientconfig = fℓ
-    result = DiffResults.GradientResult(_anyargument(ℓ)) # allocate a new result
+    result = DiffResults.GradientResult(_vectorargument(ℓ)) # allocate a new result
     result = ForwardDiff.gradient!(result, _value_closure(ℓ), x, gradientconfig)
     ValueGradient(DiffResults.value(result), DiffResults.gradient(result))
 end

src/AD_ReverseDiff.jl

@@ -0,0 +1,18 @@
+import .ReverseDiff
+
+struct ReverseDiffLogDensity{L, C} <: ADGradientWrapper
+    ℓ::L
+    gradientconfig::C
+end
+
+function ADgradient(::Val{:ReverseDiff}, ℓ::AbstractLogDensityProblem)
+    cfg = ReverseDiff.GradientConfig(zeros(dimension(ℓ)))
+    ReverseDiffLogDensity(ℓ, cfg)
+end
+
+function logdensity(::Type{ValueGradient}, fℓ::ReverseDiffLogDensity, x::RealVector)
+    @unpack ℓ, gradientconfig = fℓ
+    result = DiffResults.GradientResult(_vectorargument(ℓ)) # allocate a new result
+    result = ReverseDiff.gradient!(result, _value_closure(ℓ), x, gradientconfig)
+    ValueGradient(DiffResults.value(result), DiffResults.gradient(result))
+end

src/LogDensityProblems.jl

@@ -224,19 +224,32 @@ package being loaded.
 ADgradient(kind::Symbol, P::AbstractLogDensityProblem; kwargs...) =
     ADgradient(Val{kind}(), P; kwargs...)
 
+function ADgradient(v::Val{kind}, P::AbstractLogDensityProblem; kwargs...) where kind
+    @info "Don't know how to AD with $(kind), consider `import $(kind)` if there is such a package."
+    throw(MethodError(ADgradient, (v, P)))
+end
+
 """
 $(SIGNATURES)
 
-
+A closure for the value of the log density.
 """
 @inline _value_closure(ℓ) = x -> logdensity(Value, ℓ, x).value
 
+"""
+$(SIGNATURES)
+
+Make a vector argument for transformation `ℓ` using a Float64 vector.
+"""
+@inline _vectorargument(ℓ) = zeros(dimension(ℓ))
+
 
 # wrappers - specific
 
 function __init__()
     @require ForwardDiff="f6369f11-7733-5829-9624-2563aa707210" include("AD_ForwardDiff.jl")
     @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" include("AD_Flux.jl")
+    @require ReverseDiff="37e2e3b7-166d-5795-8a7a-e32c996b4267" include("AD_ReverseDiff.jl")
 end
 
 

test/runtests.jl

@@ -3,8 +3,7 @@ using LogDensityProblems: Value, ValueGradient
 using Test
 
 using Distributions
-import ForwardDiff
-import Flux
+import ForwardDiff, Flux, ReverseDiff
 using Parameters: @unpack
 using DocStringExtensions: SIGNATURES
 using TransformVariables
@@ -131,6 +130,15 @@ end
     @test logdensity(ValueGradient, ∇ℓ, x) ≅ ValueGradient(f(x), -6 .* x)
 end
 
+@testset "AD via ReverseDiff" begin
+    f(x) = -3*abs2(x[1])
+    ℓ = TransformedLogDensity(as(Array, asℝ, 1), f)
+    ∇ℓ = ADgradient(:ReverseDiff, ℓ)
+    x = randn(1)
+    @test logdensity(Value, ℓ, x) ≅ logdensity(Value, ∇ℓ, x)
+    @test logdensity(ValueGradient, ∇ℓ, x) ≅ ValueGradient(f(x), -6 .* x)
+end
+
 @testset "@iffinite" begin
     flag = [0]
     f(x) = (y = LogDensityProblems.@iffinite x; flag[1] += 1; y)
@@ -151,3 +159,9 @@ end
     @test logdensity(ValueGradient, ∇P, [1.0]) ≅ ValueGradient(-1.0, [-2.0])
     @test logdensity(ValueGradient, ∇P, [-1.0]) ≅ ValueGradient(-Inf, randn(1))
 end
+
+@testset "ADgradient missing method" begin
+    msg = "Don't know how to AD with Foo, consider `import Foo` if there is such a package."
+    P = TransformedLogDensity(as(Array, 1), x -> sum(abs2, x))
+    @test_logs((:info, msg), @test_throws(MethodError, ADgradient(:Foo, P)))
+end