Remove the Bijectors extension

HISTORY.md

@@ -1,3 +1,14 @@
+# Release 0.7
+
+## Removal of special treatment to `Bijectors.TransformedDistribution`
+
+Previously, `KLMinRepGradDescent`, `KLMinRepGradProxDescent`, `KLMinScoreGradDescent` only required the support of the target log-density problem to match that of `q`.
+This was implemented by giving a special treatment to `q <: Bijectors.TransformedDistribution` through the `Bijectors` extension.
+This, however, resulted in a multiplicative complexity in maintaining the relevant bits.
+Since this is not the only way to deal with constrained supports, `Bijectors` extension is now removed.
+In addition, `KLMinRepGradDescent`, `KLMinRepGradProxDescent`, `KLMinScoreGradDescent` now expect an unconstrained target log-density problem.
+Instead, a tutorial has been added to the documentation on how to deal with a target log-density problem with constrained support.
+
 # Release 0.6
 
 ## New Algorithms

Project.toml

@@ -1,6 +1,6 @@
 name = "AdvancedVI"
 uuid = "b5ca4192-6429-45e5-a2d9-87aec30a685c"
-version = "0.6"
+version = "0.7"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -20,21 +20,18 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [weakdeps]
-Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 
 [extensions]
-AdvancedVIBijectorsExt = ["Bijectors", "Optimisers"]
 AdvancedVIEnzymeExt = ["Enzyme", "ChainRulesCore"]
 AdvancedVIMooncakeExt = ["Mooncake", "ChainRulesCore"]
 AdvancedVIReverseDiffExt = ["ReverseDiff", "ChainRulesCore"]
 
 [compat]
 ADTypes = "1"
 Accessors = "0.1"
-Bijectors = "0.13, 0.14, 0.15"
 ChainRulesCore = "1"
 DiffResults = "1"
 DifferentiationInterface = "0.6, 0.7"
@@ -54,7 +51,6 @@ StatsBase = "0.32, 0.33, 0.34"
 julia = "1.10, 1.11.2"
 
 [extras]
-Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

bench/Project.toml

@@ -2,7 +2,6 @@
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AdvancedVI = "b5ca4192-6429-45e5-a2d9-87aec30a685c"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -20,9 +19,8 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ADTypes = "1"
-AdvancedVI = "0.6"
+AdvancedVI = "0.7"
 BenchmarkTools = "1"
-Bijectors = "0.13, 0.14, 0.15"
 Distributions = "0.25.111"
 DistributionsAD = "0.6"
 Enzyme = "0.13.7"

bench/benchmarks.jl

@@ -1,7 +1,6 @@
 using ADTypes
 using AdvancedVI
 using BenchmarkTools
-using Bijectors
 using Distributions
 using DistributionsAD
 using Enzyme, ForwardDiff, ReverseDiff, Zygote, Mooncake
@@ -17,8 +16,34 @@ BLAS.set_num_threads(min(4, Threads.nthreads()))
 @info sprint(versioninfo)
 @info "BLAS threads: $(BLAS.get_num_threads())"
 
-include("normallognormal.jl")
-include("unconstrdist.jl")
+struct Dist{D<:ContinuousMultivariateDistribution}
+    dist::D
+end
+
+function LogDensityProblems.logdensity(model::Dist, x)
+    return logpdf(model.dist, x)
+end
+
+function LogDensityProblems.logdensity_and_gradient(model::Dist, θ)
+    return (
+        LogDensityProblems.logdensity(model, θ),
+        ForwardDiff.gradient(Base.Fix1(LogDensityProblems.logdensity, model), θ),
+    )
+end
+
+function LogDensityProblems.dimension(model::Dist)
+    return length(model.dist)
+end
+
+function LogDensityProblems.capabilities(::Type{<:Dist})
+    return LogDensityProblems.LogDensityOrder{0}()
+end
+
+function normal(; n_dims=10, realtype=Float64)
+    μ = fill(realtype(5), n_dims)
+    Σ = Diagonal(ones(realtype, n_dims))
+    return Dist(MvNormal(μ, Σ))
+end
 
 const SUITES = BenchmarkGroup()
 
@@ -33,10 +58,7 @@ end
 begin
     T = Float64
 
-    for (probname, prob) in [
-        ("normal + bijector", normallognormal(; n_dims=10, realtype=T))
-        ("normal", normal(; n_dims=10, realtype=T))
-    ]
+    for (probname, prob) in [("normal", normal(; n_dims=10, realtype=T))]
         max_iter = 10^4
         d = LogDensityProblems.dimension(prob)
         opt = Optimisers.Adam(T(1e-3))
@@ -59,9 +81,7 @@ begin
                 ),
             ]
 
-            b = Bijectors.bijector(prob)
-            binv = inverse(b)
-            q = Bijectors.TransformedDistribution(family, binv)
+            q = family
             alg = KLMinRepGradDescent(adtype; optimizer=opt, entropy, operator=ClipScale())
 
             SUITES[probname][objname][familyname][adname] = begin

docs/Project.toml

@@ -25,7 +25,7 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 [compat]
 ADTypes = "1"
 Accessors = "0.1"
-AdvancedVI = "0.6"
+AdvancedVI = "0.7"
 Bijectors = "0.13.6, 0.14, 0.15"
 DataFrames = "1"
 DifferentiationInterface = "0.7"

docs/make.jl

@@ -21,6 +21,7 @@ makedocs(;
             "Scaling to Large Datasets" => "tutorials/subsampling.md",
             "Stan Models" => "tutorials/stan.md",
             "Normalizing Flows" => "tutorials/flows.md",
+            "Dealing with Constrained Posteriors" => "tutorials/constrained.md",
         ],
         "Algorithms" => [
             "`KLMinRepGradDescent`" => "klminrepgraddescent.md",