fixes

Author: penelopeysm

src/DynamicPPL.jl

@@ -102,6 +102,7 @@ export AbstractVarInfo,
     # LogDensityFunction
     LogDensityFunction,
     # Contexts
+    contextualize,
     SamplingContext,
     DefaultContext,
     PrefixContext,

src/experimental.jl

@@ -83,7 +83,7 @@ julia> vi isa typeof(DynamicPPL.typed_varinfo(model_with_static_support()))
 true
 ```
 """
-function determine_suitable_varinfo(model::DynamicPPL.Model, only_ddpl::Bool=true)
+function determine_suitable_varinfo(model::DynamicPPL.Model; only_ddpl::Bool=true)
     # If JET.jl has been loaded, and thus `determine_varinfo` has been defined, we use that.
     return if Base.get_extension(DynamicPPL, :DynamicPPLJETExt) !== nothing
         _determine_varinfo_jet(model; only_ddpl)

src/model.jl

@@ -85,6 +85,12 @@ function Model(f, args::NamedTuple, context::AbstractContext=DefaultContext(); k
     return Model(f, args, NamedTuple(kwargs), context)
 end
 
+"""
+    contextualize(model::Model, context::AbstractContext)
+
+Return a new `Model` with the same evaluation function and other arguments, but
+with its underlying context set to `context`.
+"""
 function contextualize(model::Model, context::AbstractContext)
     return Model(model.f, model.args, model.defaults, context)
 end

test/ad.jl

@@ -110,9 +110,8 @@ using DynamicPPL: LogDensityFunction
         # Compiling the ReverseDiff tape used to fail here
         spl = Sampler(MyEmptyAlg())
         vi = VarInfo(model)
-        ldf = LogDensityFunction(
-            model, vi, SamplingContext(spl); adtype=AutoReverseDiff(; compile=true)
-        )
+        sampling_model = contextualize(model, SamplingConext(model.context))
+        ldf = LogDensityFunction(sampling_model, vi; adtype=AutoReverseDiff(; compile=true))
         @test LogDensityProblems.logdensity_and_gradient(ldf, vi[:]) isa Any
     end