Add interface functions to allow replacing the log density function and replacing AD wrapper type

Project.toml

@@ -1,7 +1,7 @@
 name = "LogDensityProblemsAD"
 uuid = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 authors = ["Tamás K. Papp <[email protected]>"]
-version = "1.9.0"
+version = "2.0.0"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

ext/LogDensityProblemsADReverseDiffExt.jl

@@ -49,6 +49,18 @@ function ADgradient(::Val{:ReverseDiff}, ℓ;
                     x::Union{Nothing,AbstractVector}=nothing)
     ReverseDiffLogDensity(ℓ, _compiledtape(ℓ, compile, x))
 end
+function ADgradient(::Val{:ReverseDiff}, ∇ℓ::ADGradientWrapper;
+                    compile::Union{Val{true},Val{false}}=Val(false),
+                    x::Union{Nothing,AbstractVector}=nothing)
+    ADgradient(Val{:ReverseDiff}, ∇ℓ.ℓ; compile=compile, x=x)
+end
+
+function LogDensityProblemsAD.replace_ℓ(∇ℓ::ReverseDiffLogDensity{L,C}, new_ℓ) 
+    ReverseDiffLogDensity(new_ℓ, _compiledtape(new_ℓ, Val(true), nothing))
+end
+function LogDensityProblemsAD.replace_ℓ(∇ℓ::ReverseDiffLogDensity{L,Nothing}, new_ℓ) 
+    ReverseDiffLogDensity(new_ℓ, _compiledtape(new_ℓ, Val(false), nothing))
+end
 
 _compiledtape(ℓ, compile, x) = nothing
 _compiledtape(ℓ, ::Val{true}, ::Nothing) = _compiledtape(ℓ, Val(true), zeros(dimension(ℓ)))

src/LogDensityProblemsAD.jl

@@ -3,7 +3,7 @@ Automatic differentiation backends for LogDensityProblems.
 """
 module LogDensityProblemsAD
 
-export ADgradient
+export ADgradient, replace_ℓ
 
 using DocStringExtensions: SIGNATURES
 import LogDensityProblems: logdensity, logdensity_and_gradient, capabilities, dimension
@@ -38,6 +38,16 @@ Base.copy(x::ADGradientWrapper) = x # no-op, except for ForwardDiff
 """
 $(SIGNATURES)
 
+Replace the log density in an `ADGradientWrapper`.
+"""
+function replace_ℓ(∇ℓ::ADGradientWrapper, new_ℓ) 
+    @info "$(typeof(∇ℓ)) not supported for replacement"
+    throw(MethodError(replace_ℓ, (∇ℓ, new_ℓ)))
+end
+
+"""
+$(SIGNATURES)
+
 Wrap `P` using automatic differentiation to obtain a gradient.
 
 `kind` is usually a `Val` type with a symbol that refers to a package, for example