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

Make variational refs elide AD buffering

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