fix add missing files

Author: Red-Portal

src/mixedad_logdensity.jl

@@ -0,0 +1,52 @@
+
+"""
+    MixedADLogDensityProblem(problem)
+
+A `LogDensityProblem` wrapper for mixing AD frameworks.
+ Whenever the outer AD framework attempts to differentiate through `logdensity(problem)`
+the pullback calls `logdensity_and_gradient`, which invokes the inner AD framework.
+"""
+struct MixedADLogDensityProblem{Problem}
+    problem::Problem
+end
+
+function LogDensityProblems.dimension(mixedad_prob::MixedADLogDensityProblem)
+    return LogDensityProblems.dimension(mixedad_prob.problem)
+end
+
+function LogDensityProblems.logdensity(
+    mixedad_prob::MixedADLogDensityProblem, x::AbstractArray
+)
+    return LogDensityProblems.logdensity(mixedad_prob.problem, x)
+end
+
+function LogDensityProblems.logdensity_and_gradient(
+    mixedad_prob::MixedADLogDensityProblem, x::AbstractArray
+)
+    return LogDensityProblems.logdensity_and_gradient(mixedad_prob.problem, x)
+end
+
+function LogDensityProblems.logdensity_gradient_and_hessian(
+    mixedad_prob::MixedADLogDensityProblem, x::AbstractArray
+)
+    return LogDensityProblems.logdensity_gradient_and_hessian(mixedad_prob.problem, x)
+end
+
+function LogDensityProblems.capabilities(
+    ::Type{MixedADLogDensityProblem{Prob}}
+) where {Prob}
+    return LogDensityProblems.capabilities(Prob)
+end
+
+function ChainRulesCore.rrule(
+    ::typeof(LogDensityProblems.logdensity),
+    mixedad_prob::MixedADLogDensityProblem,
+    x::AbstractArray,
+)
+    ℓπ, ∇ℓπ = LogDensityProblems.logdensity_and_gradient(mixedad_prob.problem, x)
+    function logdensity_pullback(∂y)
+        ∂x = @thunk(∂y' * ∇ℓπ)
+        return NoTangent(), NoTangent(), ∂x
+    end
+    return ℓπ, logdensity_pullback
+end

test/general/mixedad_logdensity.jl

@@ -0,0 +1,40 @@
+
+struct MixedADTestModel end
+
+function LogDensityProblems.logdensity(::MixedADTestModel, θ)
+    return Float64(ℯ)
+end
+
+function LogDensityProblems.dimension(::MixedADTestModel)
+    return 3
+end
+
+function LogDensityProblems.capabilities(::Type{<:MixedADTestModel})
+    return LogDensityProblems.LogDensityOrder{2}()
+end
+
+function LogDensityProblems.logdensity_and_gradient(::MixedADTestModel, θ)
+    return (Float64(ℯ), [1.0, 2.0, 3.0])
+end
+
+function LogDensityProblems.logdensity_gradient_and_hessian(::MixedADTestModel, θ)
+    return (Float64(ℯ), [1.0, 2.0, 3.0], [1.0 1.0 1.0; 2.0 2.0 2.0; 3.0 3.0 3.0])
+end
+
+@testset "interface MixedADLogDensityProblem" begin
+    model = MixedADTestModel()
+    model_ad = AdvancedVI.MixedADLogDensityProblem(model)
+
+    d = 3
+    x = ones(Float64, d)
+
+    @test LogDensityProblems.dimension(model) == LogDensityProblems.dimension(model_ad)
+    @test LogDensityProblems.capabilities(typeof(model)) ==
+        LogDensityProblems.capabilities(typeof(model_ad))
+    @test last(LogDensityProblems.logdensity(model, x)) ≈
+        last(LogDensityProblems.logdensity(model_ad, x))
+    @test last(LogDensityProblems.logdensity_and_gradient(model, x)) ≈
+        last(LogDensityProblems.logdensity_and_gradient(model_ad, x))
+    @test last(LogDensityProblems.logdensity_gradient_and_hessian(model, x)) ≈
+        last(LogDensityProblems.logdensity_gradient_and_hessian(model_ad, x))
+end