Simplify _cache_config, drop friendly_tangents test

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: yebai

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,23 +23,19 @@ 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}})
+    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,
         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.
+# LogDensityAt is a constant w.r.t. differentiation; NoTangent avoids tangent allocation.
 Mooncake.tangent_type(::Type{<:DynamicPPL.LogDensityAt}) = NoTangent
 
 function DynamicPPL._prepare_gradient(
@@ -69,15 +61,11 @@ 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(
-    ::AutoMooncake,
+    ::Union{AutoMooncake,AutoMooncakeForward},
     prep,
     params::AbstractVector{<:Real},
     model::DynamicPPL.Model,
@@ -91,46 +79,6 @@ function DynamicPPL._value_and_gradient(
     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 adtype in (AutoMooncake(), AutoMooncakeForward())

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