Simplify Turing.Variational and fix test error (#1377)

Author: devmotion

Project.toml

@@ -20,7 +20,6 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 NamedArrays = "86f7a689-2022-50b4-a561-43c23ac3c673"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
-ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -46,7 +45,6 @@ Libtask = "0.4"
 LogDensityProblems = "^0.9, 0.10"
 MCMCChains = "4"
 NamedArrays = "0.9"
-ProgressLogging = "0.1"
 Reexport = "0.2.0"
 Requires = "0.5, 1.0"
 SpecialFunctions = "0.7.2, 0.8, 0.9, 0.10"

src/Turing.jl

@@ -15,12 +15,14 @@ using Libtask
 @reexport using Distributions, MCMCChains, Libtask, AbstractMCMC, Bijectors
 using Tracker: Tracker
 
+import AdvancedVI
 import DynamicPPL: getspace, NoDist, NamedDist
 
 const PROGRESS = Ref(true)
 function turnprogress(switch::Bool)
     @info "[Turing]: progress logging is $(switch ? "enabled" : "disabled") globally"
     PROGRESS[] = switch
+    AdvancedVI.turnprogress(switch)
 end
 
 # Random probability measures.

src/variational/VariationalInference.jl

@@ -1,23 +1,16 @@
 module Variational
 
-using ..Core, ..Utilities
-using DocStringExtensions: TYPEDEF, TYPEDFIELDS
-using Distributions, Bijectors, DynamicPPL
-using LinearAlgebra
-using ..Turing: PROGRESS, Turing
-using DynamicPPL: Model, SampleFromPrior, SampleFromUniform
-using Random: AbstractRNG
+import AdvancedVI
+import Bijectors
+import DistributionsAD
+import DynamicPPL
+import StatsBase
+import StatsFuns
 
-using ForwardDiff
-using Tracker
-
-import ..Core: getchunksize, getADbackend
-
-import AbstractMCMC
-import ProgressLogging
-
-using AdvancedVI
+import Random
 
+# Reexports
+using AdvancedVI: vi, ADVI, ELBO, elbo, TruncatedADAGrad, DecayedADAGrad
 export
     vi,
     ADVI,
@@ -34,38 +27,38 @@ use `DynamicPPL.MiniBatch` context to run the `Model` with a weight `num_total_o
 ## Notes
 - For sake of efficiency, the returned function is closes over an instance of `VarInfo`. This means that you *might* run into some weird behaviour if you call this method sequentially using different types; if that's the case, just generate a new one for each type using `make_logjoint`.
 """
-function make_logjoint(model::Model; weight = 1.0)
+function make_logjoint(model::DynamicPPL.Model; weight = 1.0)
     # setup
     ctx = DynamicPPL.MiniBatchContext(
         DynamicPPL.DefaultContext(),
         weight
     )
-    varinfo_init = Turing.VarInfo(model, ctx)
+    varinfo_init = DynamicPPL.VarInfo(model, ctx)
 
     function logπ(z)
-        varinfo = VarInfo(varinfo_init, SampleFromUniform(), z)
+        varinfo = DynamicPPL.VarInfo(varinfo_init, DynamicPPL.SampleFromUniform(), z)
         model(varinfo)
 
-        return getlogp(varinfo)
+        return DynamicPPL.getlogp(varinfo)
     end
 
     return logπ
 end
 
-function logjoint(model::Model, varinfo, z)
-    varinfo = VarInfo(varinfo, SampleFromUniform(), z)
+function logjoint(model::DynamicPPL.Model, varinfo, z)
+    varinfo = DynamicPPL.VarInfo(varinfo, DynamicPPL.SampleFromUniform(), z)
     model(varinfo)
 
-    return getlogp(varinfo)
+    return DynamicPPL.getlogp(varinfo)
 end
 
 
 # objectives
 function (elbo::ELBO)(
-    rng::AbstractRNG,
-    alg::VariationalInference,
+    rng::Random.AbstractRNG,
+    alg::AdvancedVI.VariationalInference,
     q,
-    model::Model,
+    model::DynamicPPL.Model,
     num_samples;
     weight = 1.0,
     kwargs...

src/variational/advi.jl

@@ -1,18 +1,14 @@
-using StatsFuns
-using DistributionsAD
-using Bijectors
-using Bijectors: TransformedDistribution
-using Random: AbstractRNG, GLOBAL_RNG
-import Bijectors: bijector
-
 """
-    bijector(model::Model; sym_to_ranges = Val(false))
+    bijector(model::Model[, sym2ranges = Val(false)])
 
 Returns a `Stacked <: Bijector` which maps from the support of the posterior to ℝᵈ with `d`
 denoting the dimensionality of the latent variables.
 """
-function bijector(model::Model; sym_to_ranges::Val{sym2ranges} = Val(false)) where {sym2ranges}
-    varinfo = Turing.VarInfo(model)
+function Bijectors.bijector(
+    model::DynamicPPL.Model,
+    ::Val{sym2ranges} = Val(false),
+) where {sym2ranges}
+    varinfo = DynamicPPL.VarInfo(model)
     num_params = sum([size(varinfo.metadata[sym].vals, 1)
                       for sym ∈ keys(varinfo.metadata)])
 
@@ -37,25 +33,27 @@ function bijector(model::Model; sym_to_ranges::Val{sym2ranges} = Val(false)) whe
         idx += varinfo.metadata[sym].ranges[end][end]
     end
 
-    bs = bijector.(tuple(dists...))
+    bs = Bijectors.bijector.(tuple(dists...))
 
     if sym2ranges
-        return Stacked(bs, ranges), (; collect(zip(keys(sym_lookup), values(sym_lookup)))...)
+        return (
+            Bijectors.Stacked(bs, ranges),
+            (; collect(zip(keys(sym_lookup), values(sym_lookup)))...),
+        )
     else
-        return Stacked(bs, ranges)
+        return Bijectors.Stacked(bs, ranges)
     end
 end
 
 """
-    meanfield(model::Model)
-    meanfield(rng::AbstractRNG, model::Model)
+    meanfield([rng, ]model::Model)
 
 Creates a mean-field approximation with multivariate normal as underlying distribution.
 """
-meanfield(model::Model) = meanfield(GLOBAL_RNG, model)
-function meanfield(rng::AbstractRNG, model::Model)
+meanfield(model::DynamicPPL.Model) = meanfield(Random.GLOBAL_RNG, model)
+function meanfield(rng::Random.AbstractRNG, model::DynamicPPL.Model)
     # setup
-    varinfo = Turing.VarInfo(model)
+    varinfo = DynamicPPL.VarInfo(model)
     num_params = sum([size(varinfo.metadata[sym].vals, 1)
                       for sym ∈ keys(varinfo.metadata)])
 
@@ -71,43 +69,58 @@ function meanfield(rng::AbstractRNG, model::Model)
             ranges[range_idx] = idx .+ r
             range_idx += 1
         end
-        
+
         # append!(ranges, [idx .+ r for r ∈ varinfo.metadata[sym].ranges])
         idx += varinfo.metadata[sym].ranges[end][end]
     end
 
     # initial params
     μ = randn(rng, num_params)
-    σ = softplus.(randn(rng, num_params))
+    σ = StatsFuns.softplus.(randn(rng, num_params))
 
     # construct variational posterior
-    d = TuringDiagMvNormal(μ, σ)
-    bs = inv.(bijector.(tuple(dists...)))
-    b = Stacked(bs, ranges)
+    d = DistributionsAD.TuringDiagMvNormal(μ, σ)
+    bs = inv.(Bijectors.bijector.(tuple(dists...)))
+    b = Bijectors.Stacked(bs, ranges)
 
-    return transformed(d, b)
+    return Bijectors.transformed(d, b)
 end
 
-
 # Overloading stuff from `AdvancedVI` to specialize for Turing
-AdvancedVI.update(d::TuringDiagMvNormal, μ, σ) = TuringDiagMvNormal(μ, σ)
-AdvancedVI.update(td::TransformedDistribution, θ...) = transformed(AdvancedVI.update(td.dist, θ...), td.transform)
-function AdvancedVI.update(td::TransformedDistribution{<:TuringDiagMvNormal}, θ::AbstractArray)
+function AdvancedVI.update(d::DistributionsAD.TuringDiagMvNormal, μ, σ)
+    return DistributionsAD.TuringDiagMvNormal(μ, σ)
+end
+function AdvancedVI.update(td::Bijectors.TransformedDistribution, θ...)
+    return Bijectors.transformed(AdvancedVI.update(td.dist, θ...), td.transform)
+end
+function AdvancedVI.update(
+    td::Bijectors.TransformedDistribution{<:DistributionsAD.TuringDiagMvNormal},
+    θ::AbstractArray,
+)
     μ, ω = θ[1:length(td)], θ[length(td) + 1:end]
-    return AdvancedVI.update(td, μ, softplus.(ω))
+    return AdvancedVI.update(td, μ, StatsFuns.softplus.(ω))
 end
 
-function AdvancedVI.vi(model::Model, alg::ADVI; optimizer = TruncatedADAGrad())
+function AdvancedVI.vi(
+    model::DynamicPPL.Model,
+    alg::AdvancedVI.ADVI;
+    optimizer = AdvancedVI.TruncatedADAGrad(),
+)
     q = meanfield(model)
     return AdvancedVI.vi(model, alg, q; optimizer = optimizer)
 end
 
 
-function AdvancedVI.vi(model::Model, alg::ADVI, q::TransformedDistribution{<:TuringDiagMvNormal}; optimizer = TruncatedADAGrad())
+function AdvancedVI.vi(
+    model::DynamicPPL.Model,
+    alg::AdvancedVI.ADVI,
+    q::Bijectors.TransformedDistribution{<:DistributionsAD.TuringDiagMvNormal};
+    optimizer = AdvancedVI.TruncatedADAGrad(),
+)
     @debug "Optimizing ADVI..."
     # Initial parameters for mean-field approx
-    μ, σs = params(q)
-    θ = vcat(μ, invsoftplus.(σs))
+    μ, σs = StatsBase.params(q)
+    θ = vcat(μ, StatsFuns.invsoftplus.(σs))
 
     # Optimize
     AdvancedVI.optimize!(elbo, alg, q, make_logjoint(model), θ; optimizer = optimizer)

src/variational/objectives.jl

@@ -1,5 +1,5 @@
 using Random, Test, LinearAlgebra, ForwardDiff
-using Turing.Variational: TruncatedADAGrad, DecayedADAGrad, apply!
+using AdvancedVI: TruncatedADAGrad, DecayedADAGrad, apply!
 
 function test_opt(ADPack, opt)
     θ = randn(10, 10)