Merge remote-tracking branch 'origin/v0.2-backport' into v0.2-backport

Author: torfjelde

Project.toml

@@ -1,9 +1,11 @@
 name = "AdvancedVI"
 uuid = "b5ca4192-6429-45e5-a2d9-87aec30a685c"
-version = "0.2.4"
+version = "0.2.5"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -16,22 +18,39 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
+[weakdeps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[extensions]
+AdvancedVIEnzymeExt = ["Enzyme"]
+AdvancedVIFluxExt = ["Flux"]
+AdvancedVIReverseDiffExt = ["ReverseDiff"]
+AdvancedVIZygoteExt = ["Zygote"]
+
 [compat]
 Bijectors = "0.11, 0.12, 0.13"
 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"
+Enzyme = "0.12"
+LinearAlgebra = "1.6"
 ForwardDiff = "0.10.3"
+Flux = "0.14"
 ProgressMeter = "1.0.0"
-Requires = "0.5, 1.0"
+Random = "1.6"
+Requires = "1"
+ReverseDiff = "1"
 StatsBase = "0.32, 0.33, 0.34"
 StatsFuns = "0.8, 0.9, 1"
 Tracker = "0.2.3"
+Zygote = "0.6"
 julia = "1.6"
 
 [extras]
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Pkg", "Test"]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

ext/AdvancedVIEnzymeExt.jl

@@ -0,0 +1,42 @@
+module AdvancedVIEnzymeExt
+
+if isdefined(Base, :get_extension)
+    using AdvancedVI: AdvancedVI, ADTypes, DiffResults, Distributions
+    using Enzyme: Enzyme
+else
+    using ..AdvancedVI: AdvancedVI, ADTypes, DiffResults, Distributions
+    using ..Enzyme: Enzyme
+end
+
+AdvancedVI.ADBackend(::Val{:enzyme}) = ADTypes.AutoEnzyme()
+function AdvancedVI.setadbackend(::Val{:enzyme})
+    Base.depwarn("`setadbackend` is deprecated. Please pass a `ADTypes.AbstractADType` as a keyword argument to the VI algorithm.", :setadbackend)
+    AdvancedVI.ADBACKEND[] = :enzyme
+end
+
+function AdvancedVI.grad!(
+    vo,
+    alg::AdvancedVI.VariationalInference{<:ADTypes.AutoEnzyme},
+    q,
+    model,
+    θ::AbstractVector{<:Real},
+    out::DiffResults.MutableDiffResult,
+    args...
+)
+    f(θ) =
+        if (q isa Distributions.Distribution)
+            -vo(alg, AdvancedVI.update(q, θ), model, args...)
+        else
+            -vo(alg, q(θ), model, args...)
+        end
+    # Use `Enzyme.ReverseWithPrimal` once it is released:
+    # https://github.com/EnzymeAD/Enzyme.jl/pull/598
+    y = f(θ)
+    DiffResults.value!(out, y)
+    dy = DiffResults.gradient(out)
+    fill!(dy, 0)
+    Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(θ, dy))
+    return out
+end
+
+end

ext/AdvancedVIFluxExt.jl

@@ -0,0 +1,13 @@
+module AdvancedVIFluxExt
+
+if isdefined(Base, :get_extension)
+    using AdvancedVI: AdvancedVI
+    using Flux: Flux
+else
+    using ..AdvancedVI: AdvancedVI
+    using ..Flux: Flux
+end
+
+AdvancedVI.apply!(o::Flux.Optimise.AbstractOptimiser, x, Δ) = Flux.Optimise.apply!(o, x, Δ)
+
+end

ext/AdvancedVIReverseDiffExt.jl

@@ -0,0 +1,40 @@
+module AdvancedVIReverseDiffExt
+
+if isdefined(Base, :get_extension)
+    using AdvancedVI: AdvancedVI, ADTypes, DiffResults, Distributions
+    using ReverseDiff: ReverseDiff
+else
+    using ..AdvancedVI: AdvancedVI, ADTypes, DiffResults, Distributions
+    using ..ReverseDiff: ReverseDiff
+end
+
+AdvancedVI.ADBackend(::Val{:reversediff}) = ADTypes.AutoReverseDiff()
+
+function AdvancedVI.setadbackend(::Val{:reversediff})
+    Base.depwarn("`setadbackend` is deprecated. Please pass a `ADTypes.AbstractADType` as a keyword argument to the VI algorithm.", :setadbackend)
+    AdvancedVI.ADBACKEND[] = :reversediff
+end
+
+tape(f, x) = ReverseDiff.GradientTape(f, x)
+
+function AdvancedVI.grad!(
+    vo,
+    alg::AdvancedVI.VariationalInference{<:ADTypes.AutoReverseDiff},
+    q,
+    model,
+    θ::AbstractVector{<:Real},
+    out::DiffResults.MutableDiffResult,
+    args...
+)
+    f(θ) =
+        if (q isa Distributions.Distribution)
+            -vo(alg, AdvancedVI.update(q, θ), model, args...)
+        else
+            -vo(alg, q(θ), model, args...)
+        end
+    tp = tape(f, θ)
+    ReverseDiff.gradient!(out, tp, θ)
+    return out
+end
+
+end

ext/AdvancedVIZygoteExt.jl

@@ -0,0 +1,39 @@
+module AdvancedVIZygoteExt
+
+if isdefined(Base, :get_extension)
+    using AdvancedVI: AdvancedVI, ADTypes, DiffResults, Distributions
+    using Zygote: Zygote
+else
+    using ..AdvancedVI: AdvancedVI, ADTypes, DiffResults, Distributions
+    using ..Zygote: Zygote
+end
+
+AdvancedVI.ADBackend(::Val{:zygote}) = ADTypes.AutoZygote()
+function AdvancedVI.setadbackend(::Val{:zygote})
+    Base.depwarn("`setadbackend` is deprecated. Please pass a `ADTypes.AbstractADType` as a keyword argument to the VI algorithm.", :setadbackend)
+    AdvancedVI.ADBACKEND[] = :zygote
+end
+
+function AdvancedVI.grad!(
+    vo,
+    alg::AdvancedVI.VariationalInference{<:ADTypes.AutoZygote},
+    q,
+    model,
+    θ::AbstractVector{<:Real},
+    out::DiffResults.MutableDiffResult,
+    args...
+)
+    f(θ) =
+        if (q isa Distributions.Distribution)
+            -vo(alg, AdvancedVI.update(q, θ), model, args...)
+        else
+            -vo(alg, q(θ), model, args...)
+        end
+    y, back = Zygote.pullback(f, θ)
+    dy = first(back(1.0))
+    DiffResults.value!(out, y)
+    DiffResults.gradient!(out, dy)
+    return out
+end
+
+end

src/AdvancedVI.jl

@@ -7,8 +7,11 @@ using DocStringExtensions
 
 using ProgressMeter, LinearAlgebra
 
-using ForwardDiff
-using Tracker
+using ADTypes: ADTypes
+using DiffResults: DiffResults
+
+using ForwardDiff: ForwardDiff
+using Tracker: Tracker
 
 const PROGRESS = Ref(true)
 function turnprogress(switch::Bool)
@@ -18,94 +21,6 @@ end
 
 const DEBUG = Bool(parse(Int, get(ENV, "DEBUG_ADVANCEDVI", "0")))
 
-include("ad.jl")
-include("utils.jl")
-
-using Requires
-function __init__()
-    @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" begin
-        apply!(o, x, Δ) = Flux.Optimise.apply!(o, x, Δ)
-        Flux.Optimise.apply!(o::TruncatedADAGrad, x, Δ) = apply!(o, x, Δ)
-        Flux.Optimise.apply!(o::DecayedADAGrad, x, Δ) = apply!(o, x, Δ)
-    end
-    @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" begin
-        include("compat/zygote.jl")
-        export ZygoteAD
-
-        function AdvancedVI.grad!(
-            vo,
-            alg::VariationalInference{<:AdvancedVI.ZygoteAD},
-            q,
-            model,
-            θ::AbstractVector{<:Real},
-            out::DiffResults.MutableDiffResult,
-            args...
-        )
-            f(θ) = if (q isa Distribution)
-                - vo(alg, update(q, θ), model, args...)
-            else
-                - vo(alg, q(θ), model, args...)
-            end
-            y, back = Zygote.pullback(f, θ)
-            dy = first(back(1.0))
-            DiffResults.value!(out, y)
-            DiffResults.gradient!(out, dy)
-            return out
-        end
-    end
-    @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" begin
-        include("compat/reversediff.jl")
-        export ReverseDiffAD
-
-        function AdvancedVI.grad!(
-            vo,
-            alg::VariationalInference{<:AdvancedVI.ReverseDiffAD{false}},
-            q,
-            model,
-            θ::AbstractVector{<:Real},
-            out::DiffResults.MutableDiffResult,
-            args...
-        )
-            f(θ) = if (q isa Distribution)
-                - vo(alg, update(q, θ), model, args...)
-            else
-                - vo(alg, q(θ), model, args...)
-            end
-            tp = AdvancedVI.tape(f, θ)
-            ReverseDiff.gradient!(out, tp, θ)
-            return out
-        end
-    end
-    @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" begin
-        include("compat/enzyme.jl")
-        export EnzymeAD
-
-        function AdvancedVI.grad!(
-            vo,
-            alg::VariationalInference{<:AdvancedVI.EnzymeAD},
-            q,
-            model,
-            θ::AbstractVector{<:Real},
-            out::DiffResults.MutableDiffResult,
-            args...
-        )
-            f(θ) = if (q isa Distribution)
-                - vo(alg, update(q, θ), model, args...)
-            else
-                - vo(alg, q(θ), model, args...)
-            end
-            # Use `Enzyme.ReverseWithPrimal` once it is released:
-            # https://github.com/EnzymeAD/Enzyme.jl/pull/598
-            y = f(θ)
-            DiffResults.value!(out, y)
-            dy = DiffResults.gradient(out)
-            fill!(dy, 0)
-            Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(θ, dy))
-            return out
-        end
-    end
-end
-
 export
     vi,
     ADVI,
@@ -115,10 +30,12 @@ export
     DecayedADAGrad,
     VariationalInference
 
+include("utils.jl")
+include("ad.jl")
+
 abstract type VariationalInference{AD} end
 
-getchunksize(::Type{<:VariationalInference{AD}}) where AD = getchunksize(AD)
-getADtype(::VariationalInference{AD}) where AD = AD
+getchunksize(::ADTypes.AutoForwardDiff{chunk}) where chunk = chunk === nothing ? 0 : chunk
 
 abstract type VariationalObjective end
 
@@ -129,7 +46,7 @@ const VariationalPosterior = Distribution{Multivariate, Continuous}
     grad!(vo, alg::VariationalInference, q, model::Model, θ, out, args...)
 
 Computes the gradients used in `optimize!`. Default implementation is provided for 
-`VariationalInference{AD}` where `AD` is either `ForwardDiffAD` or `TrackerAD`.
+`VariationalInference{AD}` where `AD` is either `ADTypes.AutoForwardDiff` or `ADTypes.AutoTracker`.
 This implicitly also gives a default implementation of `optimize!`.
 
 Variance reduction techniques, e.g. control variates, should be implemented in this function.
@@ -158,7 +75,7 @@ function update end
 # default implementations
 function grad!(
     vo,
-    alg::VariationalInference{<:ForwardDiffAD},
+    alg::VariationalInference{<:ADTypes.AutoForwardDiff},
     q,
     model,
     θ::AbstractVector{<:Real},
@@ -172,7 +89,7 @@ function grad!(
     end
 
     # Set chunk size and do ForwardMode.
-    chunk_size = getchunksize(typeof(alg))
+    chunk_size = getchunksize(alg.adtype)
     config = if chunk_size == 0
         ForwardDiff.GradientConfig(f, θ)
     else
@@ -183,7 +100,7 @@ end
 
 function grad!(
     vo,
-    alg::VariationalInference{<:TrackerAD},
+    alg::VariationalInference{<:ADTypes.AutoTracker},
     q,
     model,
     θ::AbstractVector{<:Real},
@@ -267,4 +184,25 @@ include("optimisers.jl")
 # VI algorithms
 include("advi.jl")
 
+if !isdefined(Base, :get_extension)
+    using Requires
+end
+
+@static if !isdefined(Base, :get_extension)
+    function __init__()
+        @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" include(
+            "../ext/AdvancedVIReverseDiffExt.jl"
+        )
+        @require Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" include(
+            "../ext/AdvancedVIZygoteExt.jl"
+        )
+        @require Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" include(
+            "../ext/AdvancedVIEnzymeExt.jl"
+        )
+        @require Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" include(
+            "../ext/AdvancedVIFluxExt.jl"
+        )
+    end
+end
+
 end # module