Merge pull request #260 from Julia-Tempering/fix-259

Make variational refs elide AD buffering

Author: miguelbiron

Project.toml

@@ -1,7 +1,7 @@
 name = "Pigeons"
 uuid = "0eb8d820-af6a-4919-95ae-11206f830c31"
 authors = ["Alexandre Bouchard-Côté <bouchard@stat.ubc.ca>, Nikola Surjanovic <nikola.surjanovic@stat.ubc.ca>, Paul Tiede <ptiede91@gmail.com>, Trevor Campbell, Miguel Biron-Lattes, Saifuddin Syed"]
-version = "0.4.3"
+version = "0.4.4"
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"

src/explorers/BufferedAD.jl

@@ -133,7 +133,7 @@ future.
     exploration step.
 """
 function get_buffer(a::Augmentation{<:Dict{Symbol, BufferedAD}}, key::Symbol, args...)
-    dict = a.contents 
+    dict = a.contents
     if !haskey(dict, key)
         dict[key] = LogDensityProblemsAD.ADgradient(args...)
     end

src/includes.jl

@@ -80,6 +80,7 @@ include("targets/DistributionLogPotential.jl")
 include("pt/checks.jl")
 include("explorers/BufferedAD.jl")
 include("variational/GaussianReference.jl")
+include("variational/VariationalReference.jl")
 include("paths/ScaledPrecisionNormalPath.jl")
 include("targets/toy_mvn_target.jl")
 include("explorers/AAPS.jl")

src/variational/GaussianReference.jl

@@ -62,8 +62,8 @@ function LogDensityProblems.dimension(log_potential::GaussianReference)
     return length(log_potential.mean[:singleton_variable])
 end
 
-LogDensityProblemsAD.ADgradient(kind::Val, log_potential::GaussianReference, buffers::Augmentation) = 
-    BufferedAD(log_potential, buffers)
+LogDensityProblemsAD.ADgradient(kind::Val, log_potential::GaussianReference, replica::Replica) = 
+    BufferedAD(log_potential, replica.recorders.buffers)
 
 function LogDensityProblems.logdensity_and_gradient(log_potential::BufferedAD{GaussianReference}, x)
     variational = log_potential.enclosed

src/variational/VariationalReference.jl

@@ -0,0 +1,19 @@
+#=
+Methods common to all variational references
+=#
+
+# Currently implemented variational references
+const VariationalReference = Union{GaussianReference}
+
+# Elide the AD buffering system
+# Reasoning: 
+#   1. Variational refs usually have analytic gradients anyway
+#   2. It can be challenging to distinguish between the proper reference and the 
+#      variational reference in the buffering system, especially since the var ref
+#      is not activated immediately
+get_buffer(
+    ::Augmentation{<:Dict{Symbol, BufferedAD}}, 
+    ::Symbol,
+    kind,
+    log_potential::VariationalReference,
+    replica::Replica) = LogDensityProblemsAD.ADgradient(kind, log_potential, replica)

test/test_BufferedAD.jl

@@ -22,18 +22,26 @@ end
 
 function test_BufferedAD_usage(pt)
     replica = last(pt.replicas)
-
+    backend = pt.shared.explorer.default_autodiff_backend
+    
     # BufferedAD were created and stored
     @test haskey(replica.recorders.ad_buffers.contents, :target)
     @test haskey(replica.recorders.ad_buffers.contents, :reference)
 
     # ADgradient uses the stored BufferedAD 
     int_lp = Pigeons.find_log_potential(replica, pt.shared.tempering, pt.shared)
-    int_ad = ADgradient(pt.shared.explorer.default_autodiff_backend, int_lp, replica)
+    int_ad = ADgradient(backend, int_lp, replica)
     @test int_ad isa Pigeons.InterpolatedAD
     @test int_ad.ref_ad === replica.recorders.ad_buffers.contents[:reference]
     @test int_ad.target_ad === replica.recorders.ad_buffers.contents[:target]
 
+    # check that the stored BufferedAD is the correct one
+    # NB: can only check type equality because extra buffers can be different (e.g. for StanLogPotential)
+    ref_ad = ADgradient(backend, int_lp.path.ref, replica)
+    target_ad = ADgradient(backend, int_lp.path.target, replica)
+    @test typeof(int_ad.ref_ad) === typeof(ref_ad)
+    @test typeof(int_ad.target_ad) === typeof(target_ad)
+
     # target and ref share the same gradient buffer
     if int_ad.ref_ad.buffer isa DiffResults.MutableDiffResult
         @test DiffResults.gradient(int_ad.ref_ad.buffer) === DiffResults.gradient(int_ad.target_ad.buffer)
@@ -100,3 +108,19 @@ end
     end
     Pigeons.set_tape_compilation_strategy!(true) # reverse setting
 end
+
+@testset "Variational reference elides the AD augmentation" begin
+    target = Pigeons.toy_stan_unid_target(100)
+    pt = pigeons(
+        target = target,
+        variational = GaussianReference(),
+        n_chains = 5,
+        n_chains_variational = 5,
+        n_rounds = 7
+    )
+    replica = pt.replicas[end];
+    var_ref = pt.shared.tempering.variational_leg.path.ref
+    var_ref_ad = Pigeons.get_buffer(replica.recorders.ad_buffers, :reference, Val(:ForwardDiff), var_ref, replica)
+    @test var_ref_ad === ADgradient(Val(:ForwardDiff), var_ref, replica)
+    @test var_ref_ad != Pigeons.get_buffer(replica.recorders.ad_buffers, :reference, Val(:ForwardDiff), target, replica)
+end