setadtype --> LogDensityFunction

Author: penelopeysm

HISTORY.md

@@ -155,13 +155,14 @@ LogDensityProblems.logdensity_and_gradient(ldf, params)
 
 without having to construct a separate `ADgradient` object.
 
-If you prefer, you can also use `setadtype` to tack on the AD type afterwards:
+If you prefer, you can also construct a new `LogDensityFunction` with a new AD type afterwards.
+The model, varinfo, and context will be taken from the original `LogDensityFunction`:
 
 ```julia
 @model f() = ...
 
 ldf = LogDensityFunction(f())  # by default, no adtype set
-ldf_with_ad = setadtype(ldf, AutoForwardDiff())
+ldf_with_ad = LogDensityFunction(ldf, AutoForwardDiff())
 ```
 
 ## 0.34.2

src/logdensityfunction.jl

@@ -143,18 +143,18 @@ struct LogDensityFunction{
 end
 
 """
-    setadtype(f::LogDensityFunction, adtype::Union{Nothing,ADTypes.AbstractADType})
-
-Set the AD type used for evaluation of log density gradient in the given
-LogDensityFunction. This function also performs preparation of the gradient,
-and sets the `prep` field of the LogDensityFunction.
-
-If `adtype` is `nothing`, the `prep` field will be set to `nothing` as well.
+    LogDensityFunction(
+        ldf::LogDensityFunction,
+        adtype::Union{Nothing,ADTypes.AbstractADType}
+    )
 
-This function returns a new LogDensityFunction with the updated AD type, i.e. it does
-not mutate the input LogDensityFunction.
+Create a new LogDensityFunction using the model, varinfo, and context from the given
+`ldf` argument, but with the AD type set to `adtype`. To remove the AD type, pass
+`nothing` as the second argument.
 """
-function setadtype(f::LogDensityFunction, adtype::Union{Nothing,ADTypes.AbstractADType})
+function LogDensityFunction(
+    f::LogDensityFunction, adtype::Union{Nothing,ADTypes.AbstractADType}
+)
     return if adtype === f.adtype
         f  # Avoid recomputing prep if not needed
     else

test/ad.jl

@@ -39,17 +39,21 @@ using DynamicPPL: LogDensityFunction
                     # Mooncake doesn't work with several combinations of SimpleVarInfo.
                     if is_mooncake && is_1_11 && is_svi_vnv
                         # https://github.com/compintell/Mooncake.jl/issues/470
-                        @test_throws ArgumentError DynamicPPL.setadtype(ref_ldf, adtype)
+                        @test_throws ArgumentError DynamicPPL.LogDensityFunction(
+                            ref_ldf, adtype
+                        )
                     elseif is_mooncake && is_1_10 && is_svi_vnv
                         # TODO: report upstream
-                        @test_throws UndefRefError DynamicPPL.setadtype(ref_ldf, adtype)
+                        @test_throws UndefRefError DynamicPPL.LogDensityFunction(
+                            ref_ldf, adtype
+                        )
                     elseif is_mooncake && is_1_10 && is_svi_od
                         # TODO: report upstream
-                        @test_throws Mooncake.MooncakeRuleCompilationError DynamicPPL.setadtype(
+                        @test_throws Mooncake.MooncakeRuleCompilationError DynamicPPL.LogDensityFunction(
                             ref_ldf, adtype
                         )
                     else
-                        ldf = DynamicPPL.setadtype(ref_ldf, adtype)
+                        ldf = DynamicPPL.LogDensityFunction(ref_ldf, adtype)
                         logp, grad = LogDensityProblems.logdensity_and_gradient(ldf, x)
                         @test grad ≈ ref_grad
                         @test logp ≈ ref_logp