From 49596f63b080948f169f0f6e16a211fc7b4a176b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 19:24:31 +0100 Subject: [PATCH 01/11] Add AutoMooncakeForward (forward-mode AD) support - Refactor `_prepare_gradient` and `_value_and_gradient` into overridable dispatch methods so backends can bypass DI entirely - Implement AutoMooncakeForward in the Mooncake extension using Mooncake's native derivative cache and a column-by-column sweep - Force `friendly_tangents=false` in `_cache_config` for both AutoMooncake and AutoMooncakeForward to keep caches valid across calls - Declare `tangent_type(LogDensityAt) = NoTangent` so Mooncake treats the function object as a constant - Relax ADP type parameter on LogDensityFunction from `Union{Nothing,DI.GradientPrep}` to unconstrained, to accommodate custom prep objects (e.g. the NamedTuple used by AutoMooncakeForward) - Add AutoMooncakeForward to the precompile workload and test suite, including a test that friendly_tangents=true config is handled correctly Co-Authored-By: Claude Sonnet 4.6 --- ext/DynamicPPLMooncakeExt.jl | 133 ++++++++++++++++++++++++++++++++--- src/logdensityfunction.jl | 103 ++++++++++++++++++--------- test/logdensityfunction.jl | 43 ++++++++--- 3 files changed, 226 insertions(+), 53 deletions(-) diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index 9760d9f4b..c090a30da 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -1,7 +1,16 @@ module DynamicPPLMooncakeExt using DynamicPPL: DynamicPPL, is_transformed -using Mooncake: Mooncake +using Mooncake: + Mooncake, + Dual, + NoTangent, + primal, + prepare_derivative_cache, + prepare_gradient_cache, + tangent, + value_and_derivative!!, + value_and_gradient!! # These are purely optimisations (although quite significant ones sometimes, especially for # _get_range_and_linked). @@ -13,19 +22,123 @@ Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{ typeof(Base.haskey),DynamicPPL.VarInfo,DynamicPPL.VarName } -using DynamicPPL: @model, LinkAll, getlogjoint_internal, LogDensityFunction -using ADTypes: AutoMooncake -import DifferentiationInterface +using DynamicPPL: @model, LinkAll, LogDensityAt, getlogjoint_internal, LogDensityFunction +using ADTypes: AutoMooncake, AutoMooncakeForward using Distributions: Normal, InverseGamma, Beta using PrecompileTools: @setup_workload, @compile_workload + +_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing}}) = Mooncake.Config() +_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) = adtype.config +function _cache_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) + config = _config(adtype) + # `friendly_tangents=true` rewrites tangent types into named structs at tape-build time, + # which is incompatible with a reusable cache (the cached tape would be tied to the + # original tangent struct layout). Force it off so the cache stays valid across calls. + return Mooncake.Config(; + debug_mode=config.debug_mode, + silence_debug_messages=config.silence_debug_messages, + friendly_tangents=false, + ) +end + +# LogDensityAt is the function being differentiated through, not a quantity being +# differentiated with respect to. Declaring NoTangent here tells Mooncake to treat it as +# a constant, which is correct and avoids unnecessary tangent allocation. +Mooncake.tangent_type(::Type{<:DynamicPPL.LogDensityAt}) = NoTangent + +function DynamicPPL._prepare_gradient( + adtype::AutoMooncake, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + return prepare_gradient_cache(f, x; config=_cache_config(adtype)) +end + +function DynamicPPL._prepare_gradient( + adtype::AutoMooncakeForward, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + return (; + cache=prepare_derivative_cache(f, x; config=_cache_config(adtype)), + dx=similar(x), + grad=similar(x), + ) +end + +function DynamicPPL._value_and_gradient( + ::AutoMooncake, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + value, (_, grad) = value_and_gradient!!(prep, f, params; args_to_zero=(false, true)) + return value, copy(grad) +end + +function DynamicPPL._value_and_gradient( + ::AutoMooncakeForward, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + dx = prep.dx + grad = prep.grad + + if isempty(grad) + # Zero-dimensional parameter vector: evaluate primal only. Use a zero tangent so + # value_and_derivative!! returns the function value without computing any derivative. + fill!(dx, zero(eltype(dx))) + value = primal( + value_and_derivative!!(prep.cache, Dual(f, NoTangent()), Dual(params, dx)) + ) + return value, copy(grad) + end + + # Standard column-by-column forward-mode sweep: set dx to each unit vector in turn, + # compute the directional derivative, and accumulate into grad. + # Each iteration resets dx[i] to zero after use, so dx is all-zeros at loop exit. + value = zero(eltype(grad)) + @inbounds for i in eachindex(grad, dx) + dx[i] = oneunit(eltype(dx)) + dual_value = value_and_derivative!!( + prep.cache, Dual(f, NoTangent()), Dual(params, dx) + ) + value = primal(dual_value) + grad[i] = tangent(dual_value) + dx[i] = zero(eltype(dx)) + end + return value, copy(grad) +end + @setup_workload begin @compile_workload begin - for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) - @model f() = x ~ dist - ldf = LogDensityFunction( - f(), getlogjoint_internal, LinkAll(); adtype=AutoMooncake() - ) - DynamicPPL.LogDensityProblems.logdensity_and_gradient(ldf, [0.5]) + for adtype in (AutoMooncake(), AutoMooncakeForward()) + for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) + @model f() = x ~ dist + ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll(); adtype) + DynamicPPL.LogDensityProblems.logdensity_and_gradient(ldf, [0.5]) + end end end end diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl index b1ae82f2b..e259165ba 100644 --- a/src/logdensityfunction.jl +++ b/src/logdensityfunction.jl @@ -178,7 +178,10 @@ struct LogDensityFunction{ L<:AbstractTransformStrategy, F, VNT<:VarNamedTuple, - ADP<:Union{Nothing,DI.GradientPrep}, + # ADP is intentionally unconstrained: most backends store a DI.GradientPrep, but + # backends that override _prepare_gradient (e.g. AutoMooncakeForward) may store any + # prep object (e.g. a NamedTuple with cache + gradient buffers). + ADP, # type of the vector passed to logdensity functions X<:AbstractVector, AC<:AccumulatorTuple, @@ -246,12 +249,9 @@ struct LogDensityFunction{ else # Make backend-specific tweaks to the adtype adtype = DynamicPPL.tweak_adtype(adtype, model, x) - args = (model, getlogdensity, all_ranges, transform_strategy, accs) - if _use_closure(adtype) - DI.prepare_gradient(LogDensityAt(args...), adtype, x) - else - DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...) - end + DynamicPPL._prepare_gradient( + adtype, x, model, getlogdensity, all_ranges, transform_strategy, accs + ) end return new{ typeof(model), @@ -407,6 +407,55 @@ function (f::LogDensityAt)(params::AbstractVector{<:Real}) ) end +function _prepare_gradient( + adtype::ADTypes.AbstractADType, + x::AbstractVector{<:Real}, + model::Model, + getlogdensity::Any, + varname_ranges::VarNamedTuple, + transform_strategy::AbstractTransformStrategy, + accs::AccumulatorTuple, +) + args = (model, getlogdensity, varname_ranges, transform_strategy, accs) + return if _use_closure(adtype) + DI.prepare_gradient(LogDensityAt(args...), adtype, x) + else + DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...) + end +end + +function _value_and_gradient( + adtype::ADTypes.AbstractADType, + prep, + params::AbstractVector{<:Real}, + model::Model, + getlogdensity::Any, + varname_ranges::VarNamedTuple, + transform_strategy::AbstractTransformStrategy, + accs::AccumulatorTuple, +) + return if _use_closure(adtype) + DI.value_and_gradient( + LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs), + prep, + adtype, + params, + ) + else + DI.value_and_gradient( + logdensity_at, + prep, + adtype, + params, + DI.Constant(model), + DI.Constant(getlogdensity), + DI.Constant(varname_ranges), + DI.Constant(transform_strategy), + DI.Constant(accs), + ) + end +end + function LogDensityProblems.logdensity( ldf::LogDensityFunction, params::AbstractVector{<:Real} ) @@ -426,32 +475,16 @@ function LogDensityProblems.logdensity_and_gradient( # `params` has to be converted to the same vector type that was used for AD preparation, # otherwise the preparation will not be valid. params = convert(get_input_vector_type(ldf), params) - return if _use_closure(ldf.adtype) - DI.value_and_gradient( - LogDensityAt( - ldf.model, - ldf._getlogdensity, - ldf._varname_ranges, - ldf.transform_strategy, - ldf._accs, - ), - ldf._adprep, - ldf.adtype, - params, - ) - else - DI.value_and_gradient( - logdensity_at, - ldf._adprep, - ldf.adtype, - params, - DI.Constant(ldf.model), - DI.Constant(ldf._getlogdensity), - DI.Constant(ldf._varname_ranges), - DI.Constant(ldf.transform_strategy), - DI.Constant(ldf._accs), - ) - end + return DynamicPPL._value_and_gradient( + ldf.adtype, + ldf._adprep, + params, + ldf.model, + ldf._getlogdensity, + ldf._varname_ranges, + ldf.transform_strategy, + ldf._accs, + ) end function LogDensityProblems.capabilities(::Type{<:LogDensityFunction{M,Nothing}}) where {M} @@ -508,6 +541,10 @@ By default, this function returns `false`, i.e. the constant approach will be us # closure (see link in the docstring). _use_closure(::ADTypes.AutoForwardDiff) = false _use_closure(::ADTypes.AutoMooncake) = false +# AutoMooncakeForward overrides _prepare_gradient/_value_and_gradient in the Mooncake +# extension and bypasses DI entirely, so this value is never reached when Mooncake is +# loaded. It is a defensive fallback for the (unlikely) case where AutoMooncakeForward is +# used without the extension. _use_closure(::ADTypes.AutoMooncakeForward) = false # For ReverseDiff, with the compiled tape, you _must_ use a closure because otherwise with # DI.Constant arguments the tape will always be recompiled upon each call to diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index b7f80a07c..cf06aa7f0 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -177,12 +177,12 @@ end struct ErrorAccumulatorException <: Exception end struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR - DynamicPPL.accumulate_assume!!( - ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any - ) = throw(ErrorAccumulatorException()) - DynamicPPL.accumulate_observe!!( - ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any - ) = throw(ErrorAccumulatorException()) + DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw( + ErrorAccumulatorException() + ) + DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw( + ErrorAccumulatorException() + ) DynamicPPL.reset(ea::ErrorAccumulator) = ea Base.copy(ea::ErrorAccumulator) = ea # Construct an LDF @@ -445,6 +445,7 @@ end AutoReverseDiff(; compile=false), AutoReverseDiff(; compile=true), AutoMooncake(; config=nothing), + AutoMooncakeForward(; config=nothing), ] @testset "Correctness" begin @@ -495,7 +496,7 @@ end return LogDensityProblems.logdensity_and_gradient(ldf, m[:]) end - @model function scalar_matrix_model(::Type{T}=Float64) where {T<:Real} + @model function scalar_matrix_model((::Type{T})=Float64) where {T<:Real} m = Matrix{T}(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end @@ -504,14 +505,14 @@ end scalar_matrix_model, test_m, ref_adtype ) - @model function matrix_model(::Type{T}=Matrix{Float64}) where {T} + @model function matrix_model((::Type{T})=Matrix{Float64}) where {T} m = T(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end matrix_model_reference = eval_logp_and_grad(matrix_model, test_m, ref_adtype) - @model function scalar_array_model(::Type{T}=Float64) where {T<:Real} + @model function scalar_array_model((::Type{T})=Float64) where {T<:Real} m = Array{T}(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end @@ -520,7 +521,7 @@ end scalar_array_model, test_m, ref_adtype ) - @model function array_model(::Type{T}=Array{Float64}) where {T} + @model function array_model((::Type{T})=Array{Float64}) where {T} m = T(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end @@ -546,6 +547,28 @@ end @test array_model_logp_and_grad[2] ≈ array_model_reference[2] end end + + @testset "Mooncake friendly_tangents" begin + @model function f() + x ~ Normal() + return y ~ Normal(x) + end + + params = randn(2) + ref_logp, ref_grad = LogDensityProblems.logdensity_and_gradient( + LogDensityFunction(f(); adtype=ref_adtype), params + ) + + for adtype in ( + AutoMooncake(; config=Mooncake.Config(; friendly_tangents=true)), + AutoMooncakeForward(; config=Mooncake.Config(; friendly_tangents=true)), + ) + ldf = LogDensityFunction(f(); adtype) + logp, grad = LogDensityProblems.logdensity_and_gradient(ldf, params) + @test logp ≈ ref_logp + @test grad ≈ ref_grad + end + end end end From f9d58e8755cc9175658d1d2659b79270c0a4d120 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Sun, 12 Apr 2026 19:29:06 +0100 Subject: [PATCH 02/11] Update test/logdensityfunction.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/logdensityfunction.jl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index cf06aa7f0..fffb1be54 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -177,12 +177,12 @@ end struct ErrorAccumulatorException <: Exception end struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR - DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw( - ErrorAccumulatorException() - ) - DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw( - ErrorAccumulatorException() - ) + DynamicPPL.accumulate_assume!!( + ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any + ) = throw(ErrorAccumulatorException()) + DynamicPPL.accumulate_observe!!( + ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any + ) = throw(ErrorAccumulatorException()) DynamicPPL.reset(ea::ErrorAccumulator) = ea Base.copy(ea::ErrorAccumulator) = ea # Construct an LDF From 2a88f94b57b3e8eab49372686f5933d327dbfeed Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 19:34:30 +0100 Subject: [PATCH 03/11] Simplify _cache_config, drop friendly_tangents test Co-Authored-By: Claude Sonnet 4.6 --- ext/DynamicPPLMooncakeExt.jl | 51 +++++------------------------------- src/logdensityfunction.jl | 9 +------ test/logdensityfunction.jl | 34 +++++------------------- 3 files changed, 14 insertions(+), 80 deletions(-) diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index c090a30da..6afd9a52a 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -3,13 +3,9 @@ module DynamicPPLMooncakeExt using DynamicPPL: DynamicPPL, is_transformed using Mooncake: Mooncake, - Dual, NoTangent, - primal, prepare_derivative_cache, prepare_gradient_cache, - tangent, - value_and_derivative!!, value_and_gradient!! # These are purely optimisations (although quite significant ones sometimes, especially for @@ -27,13 +23,11 @@ using ADTypes: AutoMooncake, AutoMooncakeForward using Distributions: Normal, InverseGamma, Beta using PrecompileTools: @setup_workload, @compile_workload -_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing}}) = Mooncake.Config() -_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) = adtype.config +function _cache_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing}}) + return Mooncake.Config(; friendly_tangents=false) +end function _cache_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) - config = _config(adtype) - # `friendly_tangents=true` rewrites tangent types into named structs at tape-build time, - # which is incompatible with a reusable cache (the cached tape would be tied to the - # original tangent struct layout). Force it off so the cache stays valid across calls. + config = adtype.config return Mooncake.Config(; debug_mode=config.debug_mode, silence_debug_messages=config.silence_debug_messages, @@ -41,9 +35,7 @@ function _cache_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) ) end -# LogDensityAt is the function being differentiated through, not a quantity being -# differentiated with respect to. Declaring NoTangent here tells Mooncake to treat it as -# a constant, which is correct and avoids unnecessary tangent allocation. +# LogDensityAt is a constant w.r.t. differentiation; NoTangent avoids tangent allocation. Mooncake.tangent_type(::Type{<:DynamicPPL.LogDensityAt}) = NoTangent function DynamicPPL._prepare_gradient( @@ -69,11 +61,7 @@ function DynamicPPL._prepare_gradient( accs::DynamicPPL.AccumulatorTuple, ) f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) - return (; - cache=prepare_derivative_cache(f, x; config=_cache_config(adtype)), - dx=similar(x), - grad=similar(x), - ) + return prepare_derivative_cache(f, x; config=_cache_config(adtype)) end function DynamicPPL._value_and_gradient( @@ -102,32 +90,7 @@ function DynamicPPL._value_and_gradient( accs::DynamicPPL.AccumulatorTuple, ) f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) - dx = prep.dx - grad = prep.grad - - if isempty(grad) - # Zero-dimensional parameter vector: evaluate primal only. Use a zero tangent so - # value_and_derivative!! returns the function value without computing any derivative. - fill!(dx, zero(eltype(dx))) - value = primal( - value_and_derivative!!(prep.cache, Dual(f, NoTangent()), Dual(params, dx)) - ) - return value, copy(grad) - end - - # Standard column-by-column forward-mode sweep: set dx to each unit vector in turn, - # compute the directional derivative, and accumulate into grad. - # Each iteration resets dx[i] to zero after use, so dx is all-zeros at loop exit. - value = zero(eltype(grad)) - @inbounds for i in eachindex(grad, dx) - dx[i] = oneunit(eltype(dx)) - dual_value = value_and_derivative!!( - prep.cache, Dual(f, NoTangent()), Dual(params, dx) - ) - value = primal(dual_value) - grad[i] = tangent(dual_value) - dx[i] = zero(eltype(dx)) - end + value, grad = value_and_gradient!!(prep, f, params) return value, copy(grad) end diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl index e259165ba..a858f7e07 100644 --- a/src/logdensityfunction.jl +++ b/src/logdensityfunction.jl @@ -178,10 +178,7 @@ struct LogDensityFunction{ L<:AbstractTransformStrategy, F, VNT<:VarNamedTuple, - # ADP is intentionally unconstrained: most backends store a DI.GradientPrep, but - # backends that override _prepare_gradient (e.g. AutoMooncakeForward) may store any - # prep object (e.g. a NamedTuple with cache + gradient buffers). - ADP, + ADP, # unconstrained: backends may store any prep object via _prepare_gradient # type of the vector passed to logdensity functions X<:AbstractVector, AC<:AccumulatorTuple, @@ -541,10 +538,6 @@ By default, this function returns `false`, i.e. the constant approach will be us # closure (see link in the docstring). _use_closure(::ADTypes.AutoForwardDiff) = false _use_closure(::ADTypes.AutoMooncake) = false -# AutoMooncakeForward overrides _prepare_gradient/_value_and_gradient in the Mooncake -# extension and bypasses DI entirely, so this value is never reached when Mooncake is -# loaded. It is a defensive fallback for the (unlikely) case where AutoMooncakeForward is -# used without the extension. _use_closure(::ADTypes.AutoMooncakeForward) = false # For ReverseDiff, with the compiled tape, you _must_ use a closure because otherwise with # DI.Constant arguments the tape will always be recompiled upon each call to diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index fffb1be54..f0d01cad1 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -177,12 +177,12 @@ end struct ErrorAccumulatorException <: Exception end struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR - DynamicPPL.accumulate_assume!!( - ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any - ) = throw(ErrorAccumulatorException()) - DynamicPPL.accumulate_observe!!( - ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any - ) = throw(ErrorAccumulatorException()) + DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw( + ErrorAccumulatorException() + ) + DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw( + ErrorAccumulatorException() + ) DynamicPPL.reset(ea::ErrorAccumulator) = ea Base.copy(ea::ErrorAccumulator) = ea # Construct an LDF @@ -547,28 +547,6 @@ end @test array_model_logp_and_grad[2] ≈ array_model_reference[2] end end - - @testset "Mooncake friendly_tangents" begin - @model function f() - x ~ Normal() - return y ~ Normal(x) - end - - params = randn(2) - ref_logp, ref_grad = LogDensityProblems.logdensity_and_gradient( - LogDensityFunction(f(); adtype=ref_adtype), params - ) - - for adtype in ( - AutoMooncake(; config=Mooncake.Config(; friendly_tangents=true)), - AutoMooncakeForward(; config=Mooncake.Config(; friendly_tangents=true)), - ) - ldf = LogDensityFunction(f(); adtype) - logp, grad = LogDensityProblems.logdensity_and_gradient(ldf, params) - @test logp ≈ ref_logp - @test grad ≈ ref_grad - end - end end end From 31cc13bad0b7b0f84e14d673924eb7114ff69928 Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Sun, 12 Apr 2026 19:51:36 +0100 Subject: [PATCH 04/11] Update test/logdensityfunction.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- ext/DynamicPPLMooncakeExt.jl | 2 +- test/logdensityfunction.jl | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index 6afd9a52a..a19ef85b9 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -96,7 +96,7 @@ end @setup_workload begin @compile_workload begin - for adtype in (AutoMooncake(), AutoMooncakeForward()) + for adtype in (AutoMooncake(),) for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) @model f() = x ~ dist ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll(); adtype) diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index f0d01cad1..3e8ce6483 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -177,12 +177,12 @@ end struct ErrorAccumulatorException <: Exception end struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR - DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw( - ErrorAccumulatorException() - ) - DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw( - ErrorAccumulatorException() - ) + DynamicPPL.accumulate_assume!!( + ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any + ) = throw(ErrorAccumulatorException()) + DynamicPPL.accumulate_observe!!( + ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any + ) = throw(ErrorAccumulatorException()) DynamicPPL.reset(ea::ErrorAccumulator) = ea Base.copy(ea::ErrorAccumulator) = ea # Construct an LDF From 0cb79f841e0a01692220864386ad1828ceaf213c Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 20:59:39 +0100 Subject: [PATCH 05/11] Use native AD APIs for ForwardDiff, Enzyme, and Mooncake; make DI optional - Move DifferentiationInterface to [weakdeps]; add DynamicPPLDifferentiationInterfaceExt as fallback for backends without native implementations - Add native ForwardDiff gradient via GradientConfig (DynamicPPLForwardDiffExt) - Add native Enzyme gradient via autodiff(ReverseWithPrimal, ...) (new DynamicPPLEnzymeExt) - Keep native Mooncake reverse/forward gradient (DynamicPPLMooncakeExt) - Add Enzyme to test env; drop DI from test env Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 9 ++- ext/DynamicPPLDifferentiationInterfaceExt.jl | 64 ++++++++++++++++++++ ext/DynamicPPLEnzymeExt.jl | 45 ++++++++++++++ ext/DynamicPPLForwardDiffExt.jl | 47 ++++++++++++-- ext/DynamicPPLMooncakeExt.jl | 18 +++++- src/logdensityfunction.jl | 52 +--------------- src/test_utils/ad.jl | 1 - test/Project.toml | 8 +-- test/logdensityfunction.jl | 15 ++--- 9 files changed, 189 insertions(+), 70 deletions(-) create mode 100644 ext/DynamicPPLDifferentiationInterfaceExt.jl create mode 100644 ext/DynamicPPLEnzymeExt.jl diff --git a/Project.toml b/Project.toml index a0b9cf0b0..57536d248 100644 --- a/Project.toml +++ b/Project.toml @@ -12,7 +12,6 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -29,6 +28,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [weakdeps] +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -38,11 +39,13 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] +DynamicPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] DynamicPPLEnzymeCoreExt = ["EnzymeCore"] +DynamicPPLEnzymeExt = ["Enzyme"] DynamicPPLForwardDiffExt = ["ForwardDiff"] DynamicPPLMCMCChainsExt = ["MCMCChains"] DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"] -DynamicPPLMooncakeExt = ["Mooncake", "DifferentiationInterface"] +DynamicPPLMooncakeExt = ["Mooncake"] DynamicPPLReverseDiffExt = ["ReverseDiff"] [compat] @@ -55,9 +58,9 @@ Bijectors = "0.15.17" Chairmarks = "1.3.1" Compat = "4" ConstructionBase = "1.5.4" -DifferentiationInterface = "0.6.41, 0.7" Distributions = "0.25" DocStringExtensions = "0.9" +Enzyme = "0.13" EnzymeCore = "0.6 - 0.8" FillArrays = "1.16.0" ForwardDiff = "0.10.12, 1" diff --git a/ext/DynamicPPLDifferentiationInterfaceExt.jl b/ext/DynamicPPLDifferentiationInterfaceExt.jl new file mode 100644 index 000000000..4b4cb8607 --- /dev/null +++ b/ext/DynamicPPLDifferentiationInterfaceExt.jl @@ -0,0 +1,64 @@ +module DynamicPPLDifferentiationInterfaceExt + +import DifferentiationInterface as DI +using DynamicPPL: + DynamicPPL, + AccumulatorTuple, + LogDensityAt, + Model, + VarNamedTuple, + AbstractTransformStrategy, + _use_closure, + logdensity_at +using ADTypes: ADTypes + +function DynamicPPL._prepare_gradient( + adtype::ADTypes.AbstractADType, + x::AbstractVector{<:Real}, + model::Model, + getlogdensity::Any, + varname_ranges::VarNamedTuple, + transform_strategy::AbstractTransformStrategy, + accs::AccumulatorTuple, +) + args = (model, getlogdensity, varname_ranges, transform_strategy, accs) + return if _use_closure(adtype) + DI.prepare_gradient(LogDensityAt(args...), adtype, x) + else + DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...) + end +end + +function DynamicPPL._value_and_gradient( + adtype::ADTypes.AbstractADType, + prep, + params::AbstractVector{<:Real}, + model::Model, + getlogdensity::Any, + varname_ranges::VarNamedTuple, + transform_strategy::AbstractTransformStrategy, + accs::AccumulatorTuple, +) + return if _use_closure(adtype) + DI.value_and_gradient( + LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs), + prep, + adtype, + params, + ) + else + DI.value_and_gradient( + logdensity_at, + prep, + adtype, + params, + DI.Constant(model), + DI.Constant(getlogdensity), + DI.Constant(varname_ranges), + DI.Constant(transform_strategy), + DI.Constant(accs), + ) + end +end + +end # module diff --git a/ext/DynamicPPLEnzymeExt.jl b/ext/DynamicPPLEnzymeExt.jl new file mode 100644 index 000000000..edc5b8f5e --- /dev/null +++ b/ext/DynamicPPLEnzymeExt.jl @@ -0,0 +1,45 @@ +module DynamicPPLEnzymeExt + +using DynamicPPL: ADTypes, DynamicPPL +using Enzyme: Enzyme + +function DynamicPPL._prepare_gradient( + ::ADTypes.AutoEnzyme, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + return (; dx=similar(x)) +end + +function DynamicPPL._value_and_gradient( + ::ADTypes.AutoEnzyme, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = DynamicPPL.LogDensityAt( + model, getlogdensity, varname_ranges, transform_strategy, accs + ) + dx = prep.dx + fill!(dx, zero(eltype(dx))) + # Const(f): LogDensityAt is not being differentiated; without Const, Enzyme errors + # because it cannot prove the function argument is readonly. + # autodiff(ReverseWithPrimal, ...) returns ((), val); dx is mutated in-place. + _, val = Enzyme.autodiff( + Enzyme.ReverseWithPrimal, + Enzyme.Const(f), + Enzyme.Active, + Enzyme.Duplicated(params, dx), + ) + return val, copy(dx) +end + +end # module diff --git a/ext/DynamicPPLForwardDiffExt.jl b/ext/DynamicPPLForwardDiffExt.jl index 4b6d3fb41..88f86fa02 100644 --- a/ext/DynamicPPLForwardDiffExt.jl +++ b/ext/DynamicPPLForwardDiffExt.jl @@ -12,10 +12,6 @@ function DynamicPPL.tweak_adtype( ) where {chunk_size} # Use DynamicPPL tag to improve stack traces # https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/ - # NOTE: DifferentiationInterface disables tag checking if the - # tag inside the AutoForwardDiff type is not nothing. See - # https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/1df562180bdcc3e91c885aa5f4162a0be2ced850/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl#L338-L350. - # So we don't currently need to override ForwardDiff.checktag as well. tag = if use_dynamicppl_tag(ad) ForwardDiff.Tag(DynamicPPL.DynamicPPLTag(), eltype(params)) else @@ -32,4 +28,47 @@ function DynamicPPL.tweak_adtype( return ADTypes.AutoForwardDiff(; chunksize=ForwardDiff.chunksize(chunk), tag=tag) end +function DynamicPPL._prepare_gradient( + adtype::ADTypes.AutoForwardDiff{chunk_size}, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) where {chunk_size} + f = DynamicPPL.LogDensityAt( + model, getlogdensity, varname_ranges, transform_strategy, accs + ) + chunk = if chunk_size == 0 || chunk_size === nothing + ForwardDiff.Chunk(x) + else + ForwardDiff.Chunk(length(x), chunk_size) + end + cfg = ForwardDiff.GradientConfig(f, x, chunk, adtype.tag) + grad = similar(x) + return (; cfg, grad) +end + +function DynamicPPL._value_and_gradient( + ::ADTypes.AutoForwardDiff, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = DynamicPPL.LogDensityAt( + model, getlogdensity, varname_ranges, transform_strategy, accs + ) + # Val{false}() skips tag checking, since our DynamicPPLTag is reused across calls + # with different LogDensityAt instances. + ForwardDiff.gradient!(prep.grad, f, params, prep.cfg, Val{false}()) + # gradient!(::AbstractArray, ...) doesn't return the value, so evaluate separately. + value = f(params) + return value, copy(prep.grad) +end + end # module diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index a19ef85b9..026ae1087 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -3,9 +3,13 @@ module DynamicPPLMooncakeExt using DynamicPPL: DynamicPPL, is_transformed using Mooncake: Mooncake, + Dual, NoTangent, prepare_derivative_cache, prepare_gradient_cache, + primal, + tangent, + value_and_derivative!!, value_and_gradient!! # These are purely optimisations (although quite significant ones sometimes, especially for @@ -61,7 +65,8 @@ function DynamicPPL._prepare_gradient( accs::DynamicPPL.AccumulatorTuple, ) f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) - return prepare_derivative_cache(f, x; config=_cache_config(adtype)) + cache = prepare_derivative_cache(f, x; config=_cache_config(adtype)) + return (; cache, dx=similar(x), grad=similar(x)) end function DynamicPPL._value_and_gradient( @@ -90,7 +95,16 @@ function DynamicPPL._value_and_gradient( accs::DynamicPPL.AccumulatorTuple, ) f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) - value, grad = value_and_gradient!!(prep, f, params) + (; cache, dx, grad) = prep + value = zero(eltype(grad)) + fill!(dx, zero(eltype(dx))) + @inbounds for i in eachindex(grad, dx) + dx[i] = one(eltype(dx)) + result = value_and_derivative!!(cache, Dual(f, NoTangent()), Dual(params, dx)) + value = primal(result) + grad[i] = tangent(result) + dx[i] = zero(eltype(dx)) + end return value, copy(grad) end diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl index a858f7e07..a3c155659 100644 --- a/src/logdensityfunction.jl +++ b/src/logdensityfunction.jl @@ -23,7 +23,6 @@ using ADTypes: ADTypes using BangBang: BangBang using AbstractPPL: AbstractPPL, VarName using LogDensityProblems: LogDensityProblems -import DifferentiationInterface as DI using Random: Random """ @@ -404,54 +403,9 @@ function (f::LogDensityAt)(params::AbstractVector{<:Real}) ) end -function _prepare_gradient( - adtype::ADTypes.AbstractADType, - x::AbstractVector{<:Real}, - model::Model, - getlogdensity::Any, - varname_ranges::VarNamedTuple, - transform_strategy::AbstractTransformStrategy, - accs::AccumulatorTuple, -) - args = (model, getlogdensity, varname_ranges, transform_strategy, accs) - return if _use_closure(adtype) - DI.prepare_gradient(LogDensityAt(args...), adtype, x) - else - DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...) - end -end - -function _value_and_gradient( - adtype::ADTypes.AbstractADType, - prep, - params::AbstractVector{<:Real}, - model::Model, - getlogdensity::Any, - varname_ranges::VarNamedTuple, - transform_strategy::AbstractTransformStrategy, - accs::AccumulatorTuple, -) - return if _use_closure(adtype) - DI.value_and_gradient( - LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs), - prep, - adtype, - params, - ) - else - DI.value_and_gradient( - logdensity_at, - prep, - adtype, - params, - DI.Constant(model), - DI.Constant(getlogdensity), - DI.Constant(varname_ranges), - DI.Constant(transform_strategy), - DI.Constant(accs), - ) - end -end +# Extensible hooks: backends provide methods via package extensions. +function _prepare_gradient end +function _value_and_gradient end function LogDensityProblems.logdensity( ldf::LogDensityFunction, params::AbstractVector{<:Real} diff --git a/src/test_utils/ad.jl b/src/test_utils/ad.jl index 8c9f96491..820a5093c 100644 --- a/src/test_utils/ad.jl +++ b/src/test_utils/ad.jl @@ -2,7 +2,6 @@ module AD using ADTypes: AbstractADType, AutoForwardDiff using Chairmarks: @be -import DifferentiationInterface as DI using DocStringExtensions using DynamicPPL: DynamicPPL, diff --git a/test/Project.toml b/test/Project.toml index 94c624616..784dd3c81 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,11 +9,12 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -43,19 +44,18 @@ BangBang = "0.4" Bijectors = "0.15.17" Chairmarks = "1" Combinatorics = "1" -DifferentiationInterface = "0.6.41, 0.7" DimensionalData = "0.30" Distributions = "0.25" Documenter = "1" +Enzyme = "0.13" ForwardDiff = "0.10.12, 1" InvertedIndices = "1" LogDensityProblems = "2" MCMCChains = "7.2.1" MacroTools = "0.5.6" MarginalLogDensities = "0.4" -Mooncake = "0.4, 0.5" -OrderedCollections = "1" OffsetArrays = "1" +OrderedCollections = "1" ReverseDiff = "1" SpecialFunctions = "2.6.1" StableRNGs = "1" diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 3e8ce6483..2f4d6b32c 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -13,9 +13,10 @@ using LogDensityProblems: LogDensityProblems using Random: Xoshiro using StableRNGs: StableRNG +using Enzyme: Enzyme using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff using Mooncake: Mooncake +using ReverseDiff: ReverseDiff @testset "LogDensityFunction: constructors" begin dist = Beta(2, 2) @@ -177,12 +178,12 @@ end struct ErrorAccumulatorException <: Exception end struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR - DynamicPPL.accumulate_assume!!( - ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any - ) = throw(ErrorAccumulatorException()) - DynamicPPL.accumulate_observe!!( - ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any - ) = throw(ErrorAccumulatorException()) + DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw( + ErrorAccumulatorException() + ) + DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw( + ErrorAccumulatorException() + ) DynamicPPL.reset(ea::ErrorAccumulator) = ea Base.copy(ea::ErrorAccumulator) = ea # Construct an LDF From de0e6b455ccd2dc0b232ffd38ad57c16c3136873 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 21:27:54 +0100 Subject: [PATCH 06/11] Address review feedback: correctness, clarity, and robustness fixes - ForwardDiff: use DiffResults (via ForwardDiff.DiffResults) for single-pass value+gradient, removing the double primal evaluation - ForwardDiff: remove redundant chunk_size guard in _prepare_gradient (tweak_adtype already normalises it to a concrete positive integer) - AutoMooncakeForward: handle empty params edge case (loop doesn't execute) - Mooncake _cache_config: use Accessors.@set to preserve all Config fields when overriding friendly_tangents=false, instead of forwarding only two known fields - Mooncake @compile_workload: remove redundant single-element for-loop - EnzymeExt: document that adtype.mode is intentionally ignored (always reverse) - src/logdensityfunction.jl: add fallback error for _value_and_gradient with unknown AD backends, pointing users to ForwardDiff (the default) or DI - test/logdensityfunction.jl: revert formatter noise (accumulate_assume!!, accumulate_observe!!, ::Type{T}=... syntax) - test/Project.toml: remove accidentally-added DynamicPPL dep Co-Authored-By: Claude Sonnet 4.6 --- Project.toml | 1 + docs/Project.toml | 2 +- ext/DynamicPPLEnzymeExt.jl | 11 +++++++++-- ext/DynamicPPLForwardDiffExt.jl | 20 +++++++++----------- ext/DynamicPPLMooncakeExt.jl | 21 +++++++++------------ src/logdensityfunction.jl | 12 ++++++++++++ test/Project.toml | 3 +-- test/logdensityfunction.jl | 20 ++++++++++---------- 8 files changed, 52 insertions(+), 38 deletions(-) diff --git a/Project.toml b/Project.toml index 57536d248..d1d254386 100644 --- a/Project.toml +++ b/Project.toml @@ -58,6 +58,7 @@ Bijectors = "0.15.17" Chairmarks = "1.3.1" Compat = "4" ConstructionBase = "1.5.4" +DifferentiationInterface = "0.6.41, 0.7" Distributions = "0.25" DocStringExtensions = "0.9" Enzyme = "0.13" diff --git a/docs/Project.toml b/docs/Project.toml index 658a85a23..c310bfb51 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -37,7 +37,7 @@ FillArrays = "0.13, 1" ForwardDiff = "0.10, 1" LogDensityProblems = "2" MCMCChains = "5, 6, 7" -MarginalLogDensities = "0.4" +MarginalLogDensities = "0.4.3" OrderedCollections = "1" StableRNGs = "1" StatsFuns = "1" diff --git a/ext/DynamicPPLEnzymeExt.jl b/ext/DynamicPPLEnzymeExt.jl index edc5b8f5e..4ea0df01c 100644 --- a/ext/DynamicPPLEnzymeExt.jl +++ b/ext/DynamicPPLEnzymeExt.jl @@ -3,6 +3,11 @@ module DynamicPPLEnzymeExt using DynamicPPL: ADTypes, DynamicPPL using Enzyme: Enzyme +_enzyme_gradient_mode(::ADTypes.AutoEnzyme{Nothing}) = Enzyme.ReverseWithPrimal +function _enzyme_gradient_mode(adtype::ADTypes.AutoEnzyme) + return Enzyme.EnzymeCore.set_runtime_activity(Enzyme.ReverseWithPrimal, adtype.mode) +end + function DynamicPPL._prepare_gradient( ::ADTypes.AutoEnzyme, x::AbstractVector{<:Real}, @@ -16,7 +21,7 @@ function DynamicPPL._prepare_gradient( end function DynamicPPL._value_and_gradient( - ::ADTypes.AutoEnzyme, + adtype::ADTypes.AutoEnzyme, prep, params::AbstractVector{<:Real}, model::DynamicPPL.Model, @@ -32,9 +37,11 @@ function DynamicPPL._value_and_gradient( fill!(dx, zero(eltype(dx))) # Const(f): LogDensityAt is not being differentiated; without Const, Enzyme errors # because it cannot prove the function argument is readonly. + # We always use reverse mode to obtain the full gradient in one pass, but preserve + # runtime-activity settings from `adtype.mode` when they were requested. # autodiff(ReverseWithPrimal, ...) returns ((), val); dx is mutated in-place. _, val = Enzyme.autodiff( - Enzyme.ReverseWithPrimal, + _enzyme_gradient_mode(adtype), Enzyme.Const(f), Enzyme.Active, Enzyme.Duplicated(params, dx), diff --git a/ext/DynamicPPLForwardDiffExt.jl b/ext/DynamicPPLForwardDiffExt.jl index 88f86fa02..b825b00a6 100644 --- a/ext/DynamicPPLForwardDiffExt.jl +++ b/ext/DynamicPPLForwardDiffExt.jl @@ -2,6 +2,9 @@ module DynamicPPLForwardDiffExt using DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems using ForwardDiff +# DiffResults is a direct dependency of ForwardDiff; access it through ForwardDiff's namespace +# rather than listing it as a separate (weak)dep of DynamicPPL. +const DiffResults = ForwardDiff.DiffResults # check if the AD type already has a tag use_dynamicppl_tag(::ADTypes.AutoForwardDiff{<:Any,Nothing}) = true @@ -40,14 +43,11 @@ function DynamicPPL._prepare_gradient( f = DynamicPPL.LogDensityAt( model, getlogdensity, varname_ranges, transform_strategy, accs ) - chunk = if chunk_size == 0 || chunk_size === nothing - ForwardDiff.Chunk(x) - else - ForwardDiff.Chunk(length(x), chunk_size) - end + # chunk_size is already set to a concrete positive integer by tweak_adtype + chunk = ForwardDiff.Chunk(length(x), chunk_size) cfg = ForwardDiff.GradientConfig(f, x, chunk, adtype.tag) - grad = similar(x) - return (; cfg, grad) + result = DiffResults.GradientResult(similar(x)) + return (; cfg, result) end function DynamicPPL._value_and_gradient( @@ -65,10 +65,8 @@ function DynamicPPL._value_and_gradient( ) # Val{false}() skips tag checking, since our DynamicPPLTag is reused across calls # with different LogDensityAt instances. - ForwardDiff.gradient!(prep.grad, f, params, prep.cfg, Val{false}()) - # gradient!(::AbstractArray, ...) doesn't return the value, so evaluate separately. - value = f(params) - return value, copy(prep.grad) + ForwardDiff.gradient!(prep.result, f, params, prep.cfg, Val{false}()) + return DiffResults.value(prep.result), copy(DiffResults.gradient(prep.result)) end end # module diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index 026ae1087..4348ba011 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -24,6 +24,7 @@ Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{ using DynamicPPL: @model, LinkAll, LogDensityAt, getlogjoint_internal, LogDensityFunction using ADTypes: AutoMooncake, AutoMooncakeForward +using Accessors: Accessors using Distributions: Normal, InverseGamma, Beta using PrecompileTools: @setup_workload, @compile_workload @@ -31,12 +32,8 @@ function _cache_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing return Mooncake.Config(; friendly_tangents=false) end function _cache_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) - config = adtype.config - return Mooncake.Config(; - debug_mode=config.debug_mode, - silence_debug_messages=config.silence_debug_messages, - friendly_tangents=false, - ) + # Use Accessors to set friendly_tangents=false while preserving all other config fields. + return Accessors.@set adtype.config.friendly_tangents = false end # LogDensityAt is a constant w.r.t. differentiation; NoTangent avoids tangent allocation. @@ -96,6 +93,8 @@ function DynamicPPL._value_and_gradient( ) f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) (; cache, dx, grad) = prep + # Handle empty parameter vector: value_and_derivative!! loop won't execute. + isempty(params) && return f(params), copy(grad) value = zero(eltype(grad)) fill!(dx, zero(eltype(dx))) @inbounds for i in eachindex(grad, dx) @@ -110,12 +109,10 @@ end @setup_workload begin @compile_workload begin - for adtype in (AutoMooncake(),) - for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) - @model f() = x ~ dist - ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll(); adtype) - DynamicPPL.LogDensityProblems.logdensity_and_gradient(ldf, [0.5]) - end + for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) + @model f() = x ~ dist + ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll(); adtype=AutoMooncake()) + DynamicPPL.LogDensityProblems.logdensity_and_gradient(ldf, [0.5]) end end end diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl index a3c155659..081ef54bb 100644 --- a/src/logdensityfunction.jl +++ b/src/logdensityfunction.jl @@ -407,6 +407,18 @@ end function _prepare_gradient end function _value_and_gradient end +function _value_and_gradient(adtype::ADTypes.AbstractADType, args...) + throw( + ArgumentError( + "No gradient implementation found for AD backend $adtype. " * + "If you intended to use the default (ForwardDiff), ensure that ForwardDiff is " * + "loaded (e.g. `using ForwardDiff`). For other backends, load the corresponding " * + "package (e.g. `using Mooncake`, `using Enzyme`) or load " * + "DifferentiationInterface as a fallback.", + ), + ) +end + function LogDensityProblems.logdensity( ldf::LogDensityFunction, params::AbstractVector{<:Real} ) diff --git a/test/Project.toml b/test/Project.toml index 784dd3c81..6d0028c6d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -13,7 +13,6 @@ DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" -DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" @@ -53,7 +52,7 @@ InvertedIndices = "1" LogDensityProblems = "2" MCMCChains = "7.2.1" MacroTools = "0.5.6" -MarginalLogDensities = "0.4" +MarginalLogDensities = "0.4.3" OffsetArrays = "1" OrderedCollections = "1" ReverseDiff = "1" diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 2f4d6b32c..d2b50aea9 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -178,12 +178,12 @@ end struct ErrorAccumulatorException <: Exception end struct ErrorAccumulator <: DynamicPPL.AbstractAccumulator end DynamicPPL.accumulator_name(::ErrorAccumulator) = :ERROR - DynamicPPL.accumulate_assume!!(::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any) = throw( - ErrorAccumulatorException() - ) - DynamicPPL.accumulate_observe!!(::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any) = throw( - ErrorAccumulatorException() - ) + DynamicPPL.accumulate_assume!!( + ::ErrorAccumulator, ::Any, ::Any, ::Any, ::VarName, ::Distribution, ::Any + ) = throw(ErrorAccumulatorException()) + DynamicPPL.accumulate_observe!!( + ::ErrorAccumulator, ::Distribution, ::Any, ::Union{VarName,Nothing}, ::Any + ) = throw(ErrorAccumulatorException()) DynamicPPL.reset(ea::ErrorAccumulator) = ea Base.copy(ea::ErrorAccumulator) = ea # Construct an LDF @@ -497,7 +497,7 @@ end return LogDensityProblems.logdensity_and_gradient(ldf, m[:]) end - @model function scalar_matrix_model((::Type{T})=Float64) where {T<:Real} + @model function scalar_matrix_model(::Type{T}=Float64) where {T<:Real} m = Matrix{T}(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end @@ -506,14 +506,14 @@ end scalar_matrix_model, test_m, ref_adtype ) - @model function matrix_model((::Type{T})=Matrix{Float64}) where {T} + @model function matrix_model(::Type{T}=Matrix{Float64}) where {T} m = T(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end matrix_model_reference = eval_logp_and_grad(matrix_model, test_m, ref_adtype) - @model function scalar_array_model((::Type{T})=Float64) where {T<:Real} + @model function scalar_array_model(::Type{T}=Float64) where {T<:Real} m = Array{T}(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end @@ -522,7 +522,7 @@ end scalar_array_model, test_m, ref_adtype ) - @model function array_model((::Type{T})=Array{Float64}) where {T} + @model function array_model(::Type{T}=Array{Float64}) where {T} m = T(undef, 2, 3) return m ~ filldist(MvNormal(zeros(2), I), 3) end From 11fae894f28570b713e7f605ff98d7d735766ffe Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 22:36:11 +0100 Subject: [PATCH 07/11] Use non-closure Enzyme logdensity call --- ext/DynamicPPLEnzymeExt.jl | 18 ++++++++---------- ext/DynamicPPLMooncakeExt.jl | 4 +++- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/ext/DynamicPPLEnzymeExt.jl b/ext/DynamicPPLEnzymeExt.jl index 4ea0df01c..e753af523 100644 --- a/ext/DynamicPPLEnzymeExt.jl +++ b/ext/DynamicPPLEnzymeExt.jl @@ -1,6 +1,6 @@ module DynamicPPLEnzymeExt -using DynamicPPL: ADTypes, DynamicPPL +using DynamicPPL: ADTypes, DynamicPPL, logdensity_at using Enzyme: Enzyme _enzyme_gradient_mode(::ADTypes.AutoEnzyme{Nothing}) = Enzyme.ReverseWithPrimal @@ -30,21 +30,19 @@ function DynamicPPL._value_and_gradient( transform_strategy::DynamicPPL.AbstractTransformStrategy, accs::DynamicPPL.AccumulatorTuple, ) - f = DynamicPPL.LogDensityAt( - model, getlogdensity, varname_ranges, transform_strategy, accs - ) dx = prep.dx fill!(dx, zero(eltype(dx))) - # Const(f): LogDensityAt is not being differentiated; without Const, Enzyme errors - # because it cannot prove the function argument is readonly. - # We always use reverse mode to obtain the full gradient in one pass, but preserve - # runtime-activity settings from `adtype.mode` when they were requested. - # autodiff(ReverseWithPrimal, ...) returns ((), val); dx is mutated in-place. + # Pass the plain function plus Const arguments; Enzyme is brittle with closure-like callables. _, val = Enzyme.autodiff( _enzyme_gradient_mode(adtype), - Enzyme.Const(f), + logdensity_at, Enzyme.Active, Enzyme.Duplicated(params, dx), + Enzyme.Const(model), + Enzyme.Const(getlogdensity), + Enzyme.Const(varname_ranges), + Enzyme.Const(transform_strategy), + Enzyme.Const(accs), ) return val, copy(dx) end diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index 4348ba011..896eed327 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -111,7 +111,9 @@ end @compile_workload begin for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) @model f() = x ~ dist - ldf = LogDensityFunction(f(), getlogjoint_internal, LinkAll(); adtype=AutoMooncake()) + ldf = LogDensityFunction( + f(), getlogjoint_internal, LinkAll(); adtype=AutoMooncake() + ) DynamicPPL.LogDensityProblems.logdensity_and_gradient(ldf, [0.5]) end end From 5ac412b39fbb21c5a938a1b2da6154a687d55ab4 Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 23:03:53 +0100 Subject: [PATCH 08/11] Cap test MarginalLogDensities for Windows CI --- test/Project.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index 6d0028c6d..edc8db0ad 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -52,7 +52,8 @@ InvertedIndices = "1" LogDensityProblems = "2" MCMCChains = "7.2.1" MacroTools = "0.5.6" -MarginalLogDensities = "0.4.3" +# 0.4.5 pulls a newer SciML stack that currently breaks Windows CI. +MarginalLogDensities = "0.4.3 - 0.4.4" OffsetArrays = "1" OrderedCollections = "1" ReverseDiff = "1" From 8c6c52a42da8a719fb5f2d8e08e17544da1210ab Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 23:30:04 +0100 Subject: [PATCH 09/11] Preserve Enzyme mode and add precedence test --- ext/DynamicPPLDifferentiationInterfaceExt.jl | 1 + ext/DynamicPPLEnzymeExt.jl | 43 +++++++++++++++++--- test/Project.toml | 3 +- 3 files changed, 40 insertions(+), 7 deletions(-) diff --git a/ext/DynamicPPLDifferentiationInterfaceExt.jl b/ext/DynamicPPLDifferentiationInterfaceExt.jl index 4b4cb8607..c76fa6f0b 100644 --- a/ext/DynamicPPLDifferentiationInterfaceExt.jl +++ b/ext/DynamicPPLDifferentiationInterfaceExt.jl @@ -12,6 +12,7 @@ using DynamicPPL: logdensity_at using ADTypes: ADTypes +# Fallback only: backend-specific extensions with more specific AD types should take precedence. function DynamicPPL._prepare_gradient( adtype::ADTypes.AbstractADType, x::AbstractVector{<:Real}, diff --git a/ext/DynamicPPLEnzymeExt.jl b/ext/DynamicPPLEnzymeExt.jl index e753af523..57dd54d88 100644 --- a/ext/DynamicPPLEnzymeExt.jl +++ b/ext/DynamicPPLEnzymeExt.jl @@ -5,11 +5,19 @@ using Enzyme: Enzyme _enzyme_gradient_mode(::ADTypes.AutoEnzyme{Nothing}) = Enzyme.ReverseWithPrimal function _enzyme_gradient_mode(adtype::ADTypes.AutoEnzyme) - return Enzyme.EnzymeCore.set_runtime_activity(Enzyme.ReverseWithPrimal, adtype.mode) + return Enzyme.EnzymeCore.WithPrimal(adtype.mode) +end + +_cache_enzyme_gradient(::ADTypes.AutoEnzyme{Nothing}) = true +_cache_enzyme_gradient(::ADTypes.AutoEnzyme{<:Enzyme.EnzymeCore.ReverseMode}) = true +_cache_enzyme_gradient(::ADTypes.AutoEnzyme) = false + +function _extract_value_and_gradient(result::NamedTuple{(:derivs, :val)}) + return result.val, first(result.derivs) end function DynamicPPL._prepare_gradient( - ::ADTypes.AutoEnzyme, + adtype::ADTypes.AutoEnzyme, x::AbstractVector{<:Real}, model::DynamicPPL.Model, getlogdensity::Any, @@ -17,12 +25,12 @@ function DynamicPPL._prepare_gradient( transform_strategy::DynamicPPL.AbstractTransformStrategy, accs::DynamicPPL.AccumulatorTuple, ) - return (; dx=similar(x)) + return _cache_enzyme_gradient(adtype) ? (; dx=similar(x)) : nothing end function DynamicPPL._value_and_gradient( adtype::ADTypes.AutoEnzyme, - prep, + prep::NamedTuple{(:dx,)}, params::AbstractVector{<:Real}, model::DynamicPPL.Model, getlogdensity::Any, @@ -32,7 +40,6 @@ function DynamicPPL._value_and_gradient( ) dx = prep.dx fill!(dx, zero(eltype(dx))) - # Pass the plain function plus Const arguments; Enzyme is brittle with closure-like callables. _, val = Enzyme.autodiff( _enzyme_gradient_mode(adtype), logdensity_at, @@ -47,4 +54,30 @@ function DynamicPPL._value_and_gradient( return val, copy(dx) end +function DynamicPPL._value_and_gradient( + adtype::ADTypes.AutoEnzyme, + ::Nothing, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + # Pass the plain function plus Const arguments; Enzyme is brittle with closure-like callables. + val, dx = _extract_value_and_gradient( + Enzyme.gradient( + _enzyme_gradient_mode(adtype), + logdensity_at, + params, + Enzyme.Const(model), + Enzyme.Const(getlogdensity), + Enzyme.Const(varname_ranges), + Enzyme.Const(transform_strategy), + Enzyme.Const(accs), + ), + ) + return val, copy(dx) +end + end # module diff --git a/test/Project.toml b/test/Project.toml index edc8db0ad..6d0028c6d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -52,8 +52,7 @@ InvertedIndices = "1" LogDensityProblems = "2" MCMCChains = "7.2.1" MacroTools = "0.5.6" -# 0.4.5 pulls a newer SciML stack that currently breaks Windows CI. -MarginalLogDensities = "0.4.3 - 0.4.4" +MarginalLogDensities = "0.4.3" OffsetArrays = "1" OrderedCollections = "1" ReverseDiff = "1" From ad487d0c202161f16643457960649b7635d9a5ed Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 23:50:53 +0100 Subject: [PATCH 10/11] Handle empty Enzyme parameter vectors --- ext/DynamicPPLEnzymeExt.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ext/DynamicPPLEnzymeExt.jl b/ext/DynamicPPLEnzymeExt.jl index 57dd54d88..591c54e38 100644 --- a/ext/DynamicPPLEnzymeExt.jl +++ b/ext/DynamicPPLEnzymeExt.jl @@ -38,6 +38,10 @@ function DynamicPPL._value_and_gradient( transform_strategy::DynamicPPL.AbstractTransformStrategy, accs::DynamicPPL.AccumulatorTuple, ) + isempty(params) && return logdensity_at( + params, model, getlogdensity, varname_ranges, transform_strategy, accs + ), + copy(params) dx = prep.dx fill!(dx, zero(eltype(dx))) _, val = Enzyme.autodiff( @@ -64,6 +68,10 @@ function DynamicPPL._value_and_gradient( transform_strategy::DynamicPPL.AbstractTransformStrategy, accs::DynamicPPL.AccumulatorTuple, ) + isempty(params) && return logdensity_at( + params, model, getlogdensity, varname_ranges, transform_strategy, accs + ), + copy(params) # Pass the plain function plus Const arguments; Enzyme is brittle with closure-like callables. val, dx = _extract_value_and_gradient( Enzyme.gradient( From 8be907f517400a2443ec93bcbf27e0faac19813b Mon Sep 17 00:00:00 2001 From: Hong Ge Date: Sun, 12 Apr 2026 23:54:15 +0100 Subject: [PATCH 11/11] Restore docs MarginalLogDensities compat --- docs/Project.toml | 2 +- ext/DynamicPPLMooncakeExt.jl | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index c310bfb51..658a85a23 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -37,7 +37,7 @@ FillArrays = "0.13, 1" ForwardDiff = "0.10, 1" LogDensityProblems = "2" MCMCChains = "5, 6, 7" -MarginalLogDensities = "0.4.3" +MarginalLogDensities = "0.4" OrderedCollections = "1" StableRNGs = "1" StatsFuns = "1" diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index 896eed327..03dfd2df7 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -66,7 +66,8 @@ function DynamicPPL._prepare_gradient( return (; cache, dx=similar(x), grad=similar(x)) end -function DynamicPPL._value_and_gradient( +# Inline this hook so the `(value, grad)` result stays as an sret tuple instead of boxing. +@inline function DynamicPPL._value_and_gradient( ::AutoMooncake, prep, params::AbstractVector{<:Real}, @@ -81,7 +82,7 @@ function DynamicPPL._value_and_gradient( return value, copy(grad) end -function DynamicPPL._value_and_gradient( +@inline function DynamicPPL._value_and_gradient( ::AutoMooncakeForward, prep, params::AbstractVector{<:Real},