Merge pull request #85 from devmotion/dw/reversediff

ReverseDiff: Do not always compile tape

Author: tpapp

Project.toml

@@ -1,7 +1,7 @@
 name = "LogDensityProblems"
 uuid = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
-version = "0.11.4"
+version = "0.11.5"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

src/AD_ReverseDiff.jl

@@ -10,24 +10,49 @@ struct ReverseDiffLogDensity{L,C} <: ADGradientWrapper
 end
 
 """
-    ADgradient(:ReverseDiff, ℓ)
-    ADgradient(Val(:ReverseDiff), ℓ)
+    ADgradient(:ReverseDiff, ℓ; compile=Val(false), x=nothing)
+    ADgradient(Val(:ReverseDiff), ℓ; compile=Val(false), x=nothing)
 
 Gradient using algorithmic/automatic differentiation via ReverseDiff.
+
+If `compile isa Val{true}`, a tape of the log density computation is created upon construction of the gradient function and used in every evaluation of the gradient.
+One may provide an example input `x::AbstractVector` of the log density function.
+If `x` is `nothing` (the default), the tape is created with input `zeros(dimension(ℓ))`.
+
+By default, no tape is created.
+
+!!! note
+    Using a compiled tape can lead to significant performance improvements when the gradient of the log density
+    is evaluated multiple times (possibly for different inputs).
+    However, if the log density contains branches, use of a compiled tape can lead to silently incorrect results.
 """
-ADgradient(::Val{:ReverseDiff}, ℓ) = begin
-    f = _logdensity_closure(ℓ)
-    x = rand(dimension(ℓ)) #init random parameters
-    tape = ReverseDiff.GradientTape(f, x)
-    compiledtape = ReverseDiff.compile(tape)
-    ReverseDiffLogDensity(ℓ, compiledtape)
+function ADgradient(::Val{:ReverseDiff}, ℓ;
+                    compile::Union{Val{true},Val{false}}=Val(false), x::Union{Nothing,AbstractVector}=nothing)
+    ReverseDiffLogDensity(ℓ, _compiledtape(ℓ, compile, x))
 end
 
-Base.show(io::IO, ∇ℓ::ReverseDiffLogDensity) = print(io, "ReverseDiff AD wrapper for ", ∇ℓ.ℓ)
+_compiledtape(ℓ, compile, x) = nothing
+_compiledtape(ℓ, ::Val{true}, ::Nothing) = _compiledtape(ℓ, Val(true), zeros(dimension(ℓ)))
+function _compiledtape(ℓ, ::Val{true}, x)
+    tape = ReverseDiff.GradientTape(Base.Fix1(logdensity, ℓ), x)
+    return ReverseDiff.compile(tape)
+end
+
+function Base.show(io::IO, ∇ℓ::ReverseDiffLogDensity)
+    print(io, "ReverseDiff AD wrapper for ", ∇ℓ.ℓ, " (")
+    if ∇ℓ.compiledtape === nothing
+        print(io, "no ")
+    end
+    print(io, "compiled tape)")
+end
 
-function logdensity_and_gradient(∇ℓ::ReverseDiffLogDensity, x::AbstractVector{T}) where {T}
+function logdensity_and_gradient(∇ℓ::ReverseDiffLogDensity, x::AbstractVector)
     @unpack ℓ, compiledtape = ∇ℓ
     buffer = _diffresults_buffer(ℓ, x)
-    result = ReverseDiff.gradient!(buffer, compiledtape, x)
+    if compiledtape === nothing
+        result = ReverseDiff.gradient!(buffer, Base.Fix1(logdensity, ℓ), x)
+    else
+        result = ReverseDiff.gradient!(buffer, compiledtape, x)
+    end
     _diffresults_extract(result)
 end

test/runtests.jl

@@ -96,16 +96,42 @@ end
 ####
 
 @testset "AD via ReverseDiff" begin
-    ℓ = TestLogDensity(test_logdensity1)
-    ∇ℓ = ADgradient(:ReverseDiff, ℓ)
-    @test repr(∇ℓ) == "ReverseDiff AD wrapper for " * repr(ℓ)
-    @test dimension(∇ℓ) == 3
-    @test capabilities(∇ℓ) ≡ LogDensityOrder(1)
-    for _ in 1:100
-        x = randn(3)
-        @test @inferred(logdensity(∇ℓ, x)) ≅ test_logdensity1(x)
-        @test @inferred(logdensity_and_gradient(∇ℓ, x)) ≅
-            (test_logdensity1(x), test_gradient(x))
+    ℓ = TestLogDensity()
+
+    ∇ℓ_default = ADgradient(:ReverseDiff, ℓ)
+    ∇ℓ_nocompile = ADgradient(:ReverseDiff, ℓ; compile=Val(false))
+    for ∇ℓ in (∇ℓ_default, ∇ℓ_nocompile)
+        @test repr(∇ℓ) == "ReverseDiff AD wrapper for " * repr(ℓ) * " (no compiled tape)"
+    end
+
+    ∇ℓ_compile = ADgradient(:ReverseDiff, ℓ; compile=Val(true))
+    ∇ℓ_compile_x = ADgradient(:ReverseDiff, ℓ; compile=Val(true), x=rand(3))
+    for ∇ℓ in (∇ℓ_compile, ∇ℓ_compile_x)
+        @test repr(∇ℓ) == "ReverseDiff AD wrapper for " * repr(ℓ) * " (compiled tape)"
+    end
+
+    for ∇ℓ in (∇ℓ_default, ∇ℓ_nocompile, ∇ℓ_compile, ∇ℓ_compile_x)
+        @test dimension(∇ℓ) == 3
+        @test capabilities(∇ℓ) ≡ LogDensityOrder(1)
+
+        for _ in 1:100
+            x = rand(3)
+            @test @inferred(logdensity(∇ℓ, x)) ≅ test_logdensity(x)
+            @test @inferred(logdensity_and_gradient(∇ℓ, x)) ≅
+                (test_logdensity(x), test_gradient(x))
+
+            x = -x
+            @test @inferred(logdensity(∇ℓ, x)) ≅ test_logdensity(x)
+            if ∇ℓ.compiledtape === nothing
+                # Recompute tape => correct results
+                @test @inferred(logdensity_and_gradient(∇ℓ, x)) ≅
+                    (test_logdensity(x), zero(x))
+            else
+                # Tape not recomputed => incorrect results, uses always the same branch
+                @test @inferred(logdensity_and_gradient(∇ℓ, x)) ≅
+                    (test_logdensity1(x), test_gradient(x))
+            end
+        end
     end
 end