Use LogDensityFunction for variational inference (#1944)

* use LogDensityFunction for vi

* bump minor version

* removed undefined reference to Bijectors.setadbacked

* removed incorrect previous commit

* Update VariationalInference.jl

---------

Co-authored-by: Hong Ge <[email protected]>

Author: torfjelde

src/variational/VariationalInference.jl

@@ -1,14 +1,17 @@
 module Variational
 
-using AdvancedVI
-using Bijectors
-using DistributionsAD
-using DynamicPPL
-using StatsBase
-using StatsFuns
+using DistributionsAD: DistributionsAD
+using DynamicPPL: DynamicPPL
+using StatsBase: StatsBase
+using StatsFuns: StatsFuns
+using LogDensityProblems: LogDensityProblems
 using Distributions
 
-using Random
+using Random: Random
+
+import AdvancedVI
+import Bijectors
+
 
 # Reexports
 using AdvancedVI: vi, ADVI, ELBO, elbo, TruncatedADAGrad, DecayedADAGrad
@@ -34,26 +37,11 @@ function make_logjoint(model::DynamicPPL.Model; weight = 1.0)
         DynamicPPL.DefaultContext(),
         weight
     )
-    varinfo_init = DynamicPPL.VarInfo(model, ctx)
-
-    function logπ(z)
-        varinfo = DynamicPPL.VarInfo(varinfo_init, DynamicPPL.SampleFromUniform(), z)
-        model(varinfo)
-
-        return DynamicPPL.getlogp(varinfo)
-    end
-
-    return logπ
+    model_contextualized = DynamicPPL.contextualize(model, ctx)
+    f = DynamicPPL.LogDensityFunction(model_contextualized)
+    return Base.Fix1(LogDensityProblems.logdensity, f)
 end
 
-function logjoint(model::DynamicPPL.Model, varinfo, z)
-    varinfo = DynamicPPL.VarInfo(varinfo, DynamicPPL.SampleFromUniform(), z)
-    model(varinfo)
-
-    return DynamicPPL.getlogp(varinfo)
-end
-
-
 # objectives
 function (elbo::ELBO)(
     rng::Random.AbstractRNG,