Enable the use of structured gradients for Zygote

Project.toml

@@ -6,7 +6,6 @@ version = "0.3.0"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -39,7 +38,6 @@ ADTypes = "0.1, 0.2, 1"
 Accessors = "0.1"
 Bijectors = "0.13"
 ChainRulesCore = "1.16"
-DiffResults = "1"
 Distributions = "0.25.87"
 DocStringExtensions = "0.8, 0.9"
 Enzyme = "0.12"

ext/AdvancedVIEnzymeExt.jl

@@ -4,23 +4,22 @@ module AdvancedVIEnzymeExt
 if isdefined(Base, :get_extension)
     using Enzyme
     using AdvancedVI
-    using AdvancedVI: ADTypes, DiffResults
+    using AdvancedVI: ADTypes
 else
     using ..Enzyme
     using ..AdvancedVI
-    using ..AdvancedVI: ADTypes, DiffResults
+    using ..AdvancedVI: ADTypes
 end
 
 # Enzyme doesn't support f::Bijectors (see https://github.com/EnzymeAD/Enzyme.jl/issues/916)
-function AdvancedVI.value_and_gradient!(
-    ad::ADTypes.AutoEnzyme, f, θ::AbstractVector{T}, out::DiffResults.MutableDiffResult
+function AdvancedVI.value_and_gradient(
+    ad::ADTypes.AutoEnzyme, f, θ::AbstractVector{T}
 ) where {T<:Real}
     y = f(θ)
-    DiffResults.value!(out, y)
-    ∇θ = DiffResults.gradient(out)
+    ∇θ = similar(θ)
     fill!(∇θ, zero(T))
     Enzyme.autodiff(Enzyme.ReverseWithPrimal, f, Enzyme.Active, Enzyme.Duplicated(θ, ∇θ))
-    return out
+    ∇θ,  y
 end
 
 end

ext/AdvancedVIForwardDiffExt.jl

@@ -4,26 +4,26 @@ module AdvancedVIForwardDiffExt
 if isdefined(Base, :get_extension)
     using ForwardDiff
     using AdvancedVI
-    using AdvancedVI: ADTypes, DiffResults
+    using AdvancedVI: ADTypes
 else
     using ..ForwardDiff
     using ..AdvancedVI
-    using ..AdvancedVI: ADTypes, DiffResults
+    using ..AdvancedVI: ADTypes
 end
 
 getchunksize(::ADTypes.AutoForwardDiff{chunksize}) where {chunksize} = chunksize
 
-function AdvancedVI.value_and_gradient!(
-    ad::ADTypes.AutoForwardDiff, f, θ::AbstractVector{T}, out::DiffResults.MutableDiffResult
-) where {T<:Real}
+function AdvancedVI.value_and_gradient(
+    ad::ADTypes.AutoForwardDiff, f, θ::AbstractVector{<:Real}
+)
     chunk_size = getchunksize(ad)
     config = if isnothing(chunk_size)
         ForwardDiff.GradientConfig(f, θ)
     else
         ForwardDiff.GradientConfig(f, θ, ForwardDiff.Chunk(length(θ), chunk_size))
     end
-    ForwardDiff.gradient!(out, f, θ, config)
-    return out
+    g = ForwardDiff.gradient(f, θ, config)
+    g, f(θ)
 end
 
 end

ext/AdvancedVIReverseDiffExt.jl

@@ -3,21 +3,21 @@ module AdvancedVIReverseDiffExt
 
 if isdefined(Base, :get_extension)
     using AdvancedVI
-    using AdvancedVI: ADTypes, DiffResults
+    using AdvancedVI: ADTypes
     using ReverseDiff
 else
     using ..AdvancedVI
-    using ..AdvancedVI: ADTypes, DiffResults
+    using ..AdvancedVI: ADTypes
     using ..ReverseDiff
 end
 
 # ReverseDiff without compiled tape
-function AdvancedVI.value_and_gradient!(
-    ad::ADTypes.AutoReverseDiff, f, θ::AbstractVector{<:Real}, out::DiffResults.MutableDiffResult
+function AdvancedVI.value_and_gradient(
+    ad::ADTypes.AutoReverseDiff, f, θ::AbstractVector{<:Real}
 )
     tp = ReverseDiff.GradientTape(f, θ)
-    ReverseDiff.gradient!(out, tp, θ)
-    return out
+    g = ReverseDiff.gradient!(tp, θ)
+    g, f(θ)
 end
 
 end

ext/AdvancedVIZygoteExt.jl

@@ -3,22 +3,18 @@ module AdvancedVIZygoteExt
 
 if isdefined(Base, :get_extension)
     using AdvancedVI
-    using AdvancedVI: ADTypes, DiffResults
+    using AdvancedVI: ADTypes
     using Zygote
 else
     using ..AdvancedVI
-    using ..AdvancedVI: ADTypes, DiffResults
+    using ..AdvancedVI: ADTypes
     using ..Zygote
 end
 
-function AdvancedVI.value_and_gradient!(
-    ad::ADTypes.AutoZygote, f, θ::AbstractVector{<:Real}, out::DiffResults.MutableDiffResult
-)
+function AdvancedVI.value_and_gradient(ad::ADTypes.AutoZygote, f, θ)
     y, back = Zygote.pullback(f, θ)
     ∇θ = back(one(y))
-    DiffResults.value!(out, y)
-    DiffResults.gradient!(out, only(∇θ))
-    return out
+    only(∇θ), y
 end
 
 end

src/AdvancedVI.jl

@@ -16,7 +16,7 @@ using LinearAlgebra
 
 using LogDensityProblems
 
-using ADTypes, DiffResults
+using ADTypes
 using ChainRulesCore
 
 using FillArrays
@@ -25,17 +25,43 @@ using StatsBase
 
 # derivatives
 """
-    value_and_gradient!(ad, f, θ, out)
+    value_and_gradient(ad, f, θ)
 
-Evaluate the value and gradient of a function `f` at `θ` using the automatic differentiation backend `ad` and store the result in `out`.
+Evaluate the value and gradient of a function `f` at `θ` using the automatic differentiation backend `ad`.
 
 # Arguments
 - `ad::ADTypes.AbstractADType`: Automatic differentiation backend. 
 - `f`: Function subject to differentiation.
 - `θ`: The point to evaluate the gradient.
 - `out::DiffResults.MutableDiffResult`: Buffer to contain the output gradient and function value.
+
+# Returns
+- `grad`: Gradient of `f` evaluated on `θ`
+- `fval`: Function value `f` evaluated on `θ`
+"""
+function value_and_gradient end
+
+maybe_destructure(::ADTypes.AutoZygote, q) = (q, identity)
+
+maybe_destructure(::ADTypes.AbstractADType, q) = Optimisers.destructure(q)
+
 """
-function value_and_gradient! end
+    update_variational_params!(family_type, opt_st, params, re, grad)
+
+Update variational family according to 
+Essentially an indirection for `Optimisers.update!`.
+
+# Arguments
+- `family_type::Type`: 
+
+# Returns
+- `opt_st`: Updated optimizer state.
+- `params`: Updated params.
+"""
+function update_variational_params! end
+
+update_variational_params!(::Type, opt_st, params, re, grad) =
+    Optimisers.update!(opt_st, params, grad)
 
 # estimators
 """
@@ -51,23 +77,18 @@ If the estimator is stateful, it can implement `init` to initialize the state.
 abstract type AbstractVariationalObjective end
 
 """
-    init(rng, obj, λ, restructure)
+    init(rng, obj, params, q_init; kwargs...)
 
 Initialize a state of the variational objective `obj` given the initial variational parameters `λ`.
 This function needs to be implemented only if `obj` is stateful.
 
 # Arguments
 - `rng::Random.AbstractRNG`: Random number generator.
 - `obj::AbstractVariationalObjective`: Variational objective.
-- `λ`: Initial variational parameters.
-- `restructure`: Function that reconstructs the variational approximation from `λ`.
+- `params`: Initial variational parameters.
+- `q_init`: Initial variational distribution.
 """
-init(
-    ::Random.AbstractRNG,
-    ::AbstractVariationalObjective,
-    ::AbstractVector,
-    ::Any
-) = nothing
+init(::Random.AbstractRNG, ::AbstractVariationalObjective, ::Any, ::Any; kwargs...) = nothing
 
 """
     estimate_objective([rng,] obj, q, prob; kwargs...)
@@ -91,28 +112,26 @@ function estimate_objective end
 
 export estimate_objective
 
-
 """
-    estimate_gradient!(rng, obj, adtype, out, prob, λ, restructure, obj_state)
+    estimate_gradient(rng, obj, adtype, prob, params, restructure, obj_state; kwargs...)
 
 Estimate (possibly stochastic) gradients of the variational objective `obj` targeting `prob` with respect to the variational parameters `λ`
 
 # Arguments
 - `rng::Random.AbstractRNG`: Random number generator.
 - `obj::AbstractVariationalObjective`: Variational objective.
 - `adtype::ADTypes.AbstractADType`: Automatic differentiation backend. 
-- `out::DiffResults.MutableDiffResult`: Buffer containing the objective value and gradient estimates. 
 - `prob`: The target log-joint likelihood implementing the `LogDensityProblem` interface.
-- `λ`: Variational parameters to evaluate the gradient on.
+- `params`: Variational parameters to evaluate the gradient on.
 - `restructure`: Function that reconstructs the variational approximation from `λ`.
 - `obj_state`: Previous state of the objective.
 
 # Returns
-- `out::MutableDiffResult`: Buffer containing the objective value and gradient estimates.
+- `grad`: Gradient estimate.
 - `obj_state`: The updated state of the objective.
 - `stat::NamedTuple`: Statistics and logs generated during estimation.
 """
-function estimate_gradient! end
+function estimate_gradient end
 
 # ELBO-specific interfaces
 abstract type AbstractEntropyEstimator end

src/families/location_scale.jl

@@ -57,17 +57,17 @@ Base.eltype(::Type{<:MvLocationScale{S, D, L}}) where {S, D, L} = eltype(D)
 function StatsBase.entropy(q::MvLocationScale)
     @unpack  location, scale, dist = q
     n_dims = length(location)
-    n_dims*convert(eltype(location), entropy(dist)) + first(logabsdet(scale))
+    n_dims*convert(eltype(location), entropy(dist)) + first(logdet(scale))
 end
 
 function Distributions.logpdf(q::MvLocationScale, z::AbstractVector{<:Real})
     @unpack location, scale, dist = q
-    sum(Base.Fix1(logpdf, dist), scale \ (z - location)) - first(logabsdet(scale))
+    sum(Base.Fix1(logpdf, dist), scale \ (z - location)) - first(logdet(scale))
 end
 
 function Distributions._logpdf(q::MvLocationScale, z::AbstractVector{<:Real})
     @unpack location, scale, dist = q
-    sum(Base.Fix1(logpdf, dist), scale \ (z - location)) - first(logabsdet(scale))
+    sum(Base.Fix1(logpdf, dist), scale \ (z - location)) - first(logdet(scale))
 end
 
 function Distributions.rand(q::MvLocationScale)

src/objectives/elbo/repgradelbo.jl

@@ -91,31 +91,27 @@ function estimate_objective(
     energy + entropy
 end
 
-estimate_objective(obj::RepGradELBO, q, prob; n_samples::Int = obj.n_samples) =
+estimate_objective(obj::RepGradELBO, q, prob; n_samples::Int = obj.n_samples, kwargs...) =
     estimate_objective(Random.default_rng(), obj, q, prob; n_samples)
 
-function estimate_gradient!(
+function estimate_gradient(
     rng   ::Random.AbstractRNG,
     obj   ::RepGradELBO,
     adtype::ADTypes.AbstractADType,
-    out   ::DiffResults.MutableDiffResult,
     prob,
-    λ,
+    params,
     restructure,
-    state,
+    state;
+    kwargs...
 )
-    q_stop = restructure(λ)
-    function f(λ′)
-        q = restructure(λ′)
+    q_stop = restructure(params)
+    function f(params′)
+        q = restructure(params′)
         samples, entropy = reparam_with_entropy(rng, q, q_stop, obj.n_samples, obj.entropy)
         energy = estimate_energy_with_samples(prob, samples)
         elbo = energy + entropy
         -elbo
     end
-    value_and_gradient!(adtype, f, λ, out)
-
-    nelbo = DiffResults.value(out)
-    stat  = (elbo=-nelbo,)
-
-    out, nothing, stat
+    grad, nelbo = value_and_gradient(adtype, f, params)
+    grad, nothing, (elbo=-nelbo,)
 end