bump Bijectors.jl compat bounds and replace forward with rand_and_logjac

Author: torfjelde

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedVI"
 uuid = "b5ca4192-6429-45e5-a2d9-87aec30a685c"
-version = "0.1.6"
+version = "0.2.0"
 
 [deps]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
@@ -17,7 +17,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [compat]
-Bijectors = "0.4.0, 0.5, 0.6, 0.7, 0.8, 0.9, 0.10"
+Bijectors = "0.11, 0.12"
 Distributions = "0.21, 0.22, 0.23, 0.24, 0.25"
 DistributionsAD = "0.2, 0.3, 0.4, 0.5, 0.6"
 DocStringExtensions = "0.8, 0.9"

src/AdvancedVI.jl

@@ -19,6 +19,7 @@ end
 const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_ADVANCEDVI", "0")))
 
 include("ad.jl")
+include("utils.jl")
 
 using Requires
 function __init__()

src/advi.jl

@@ -81,8 +81,8 @@ function (elbo::ELBO)(
     #      = 𝔼[log p(x, f⁻¹(z̃)) + logabsdet(J(f⁻¹(z̃)))] + ℍ(q̃(z̃))
     #      = 𝔼[log p(x, z) - logabsdetjac(J(f(z)))] + ℍ(q̃(z̃))
 
-    # But our `forward(q)` is using f⁻¹: ℝ → supp(p(z | x)) going forward → `+ logjac`
-    _, z, logjac, _ = forward(rng, q)
+    # But our `rand_and_logjac(q)` is using f⁻¹: ℝ → supp(p(z | x)) going forward → `+ logjac`
+    z, logjac = rand_and_logjac(rng, q)
     res = (logπ(z) + logjac) / num_samples
 
     if q isa TransformedDistribution
@@ -92,7 +92,7 @@ function (elbo::ELBO)(
     end
 
     for i = 2:num_samples
-        _, z, logjac, _ = forward(rng, q)
+        z, logjac = rand_and_logjac(rng, q)
         res += (logπ(z) + logjac) / num_samples
     end
 

src/utils.jl

@@ -0,0 +1,16 @@
+using Distributions
+
+using Random: Random
+using Bijectors: Bijectors
+
+
+function rand_and_logjac(rng::Random.AbstractRNG, dist::Distribution)
+    x = rand(rng, dist)
+    return x, zero(eltype(x))
+end
+
+function rand_and_logjac(rng::Random.AbstractRNG, dist::Bijectors.TransformedDistribution)
+    x = rand(rng, dist.dist)
+    y, logjac = Bijectors.with_logabsdet_jacobian(dist.transform, x)
+    return y, logjac
+end