Fix tests

Author: penelopeysm

src/logdensityfunction.jl

@@ -4,8 +4,9 @@ import DifferentiationInterface as DI
     LogDensityFunction(
         model::Model,
         varinfo::AbstractVarInfo=VarInfo(model),
-        context::AbstractContext=DefaultContext(),
-        adtype::Union{ADTypes.AbstractADType,Nothing}=nothing)
+        context::AbstractContext=DefaultContext();
+        adtype::Union{ADTypes.AbstractADType,Nothing}=nothing
+    )
 
 A struct which contains a model, along with all the information necessary to:
 
@@ -16,8 +17,9 @@ At its most basic level, a LogDensityFunction wraps the model together with its
 the type of varinfo to be used, as well as the evaluation context. These must
 be known in order to calculate the log density (using [`DynamicPPL.evaluate!!`](@ref)).
 
-If `adtype` is provided, then this struct will also contain the adtype along with
-other information for efficient calculation of the gradient of the log density.
+If the `adtype` keyword argument is provided, then this struct will also
+store the adtype along with other information for efficient calculation of the
+gradient of the log density.
 
 `DynamicPPL.LogDensityFunction` implements the LogDensityProblems.jl interface.
 If `adtype` is nothing, then only `logdensity` is implemented. If `adtype` is a
@@ -27,6 +29,7 @@ concrete AD backend type, then `logdensity_and_gradient` is also implemented.
 $(FIELDS)
 
 # Examples
+
 ```jldoctest
 julia> using Distributions
 
@@ -62,18 +65,26 @@ julia> # This also respects the context in `model`.
 
 julia> LogDensityProblems.logdensity(f_prior, [0.0]) == logpdf(Normal(), 0.0)
 true
+
+julia> # If we also need to calculate the gradient, we can specify an AD backend.
+       import ForwardDiff, ADTypes
+
+julia> f = LogDensityFunction(model, adtype=ADTypes.AutoForwardDiff());
+
+julia> LogDensityProblems.logdensity_and_gradient(f, [0.0])
+(-2.3378770664093453, [1.0])
 ```
 """
 struct LogDensityFunction{
-    V<:AbstractVarInfo,
     M<:Model,
-    C<:Union{Nothing,AbstractContext},
+    V<:AbstractVarInfo,
+    C<:AbstractContext,
     AD<:Union{Nothing,ADTypes.AbstractADType},
 }
-    "varinfo used for evaluation"
-    varinfo::V
     "model used for evaluation"
     model::M
+    "varinfo used for evaluation"
+    varinfo::V
     "context used for evaluation; if `nothing`, `leafcontext(model.context)` will be used when applicable"
     context::C
     "AD type used for evaluation of log density gradient. If `nothing`, no gradient can be calculated"
@@ -86,7 +97,7 @@ struct LogDensityFunction{
     function LogDensityFunction(
         model::Model,
         varinfo::AbstractVarInfo=VarInfo(model),
-        context::AbstractContext=DefaultContext();
+        context::AbstractContext=leafcontext(model.context);
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing,
     )
         if adtype === nothing
@@ -112,8 +123,8 @@ struct LogDensityFunction{
             end
             with_closure = with_closure
         end
-        return new{typeof(varinfo),typeof(model),typeof(context),typeof(adtype)}(
-            varinfo, model, context, adtype, prep, with_closure
+        return new{typeof(model),typeof(varinfo),typeof(context),typeof(adtype)}(
+            model, varinfo, context, adtype, prep, with_closure
         )
     end
 end
@@ -161,21 +172,21 @@ end
 ### LogDensityProblems interface
 
 function LogDensityProblems.capabilities(
-    ::Type{<:LogDensityFunction{V,M,C,Nothing}}
-) where {V,M,C}
+    ::Type{<:LogDensityFunction{M,V,C,Nothing}}
+) where {M,V,C}
     return LogDensityProblems.LogDensityOrder{0}()
 end
 function LogDensityProblems.capabilities(
-    ::Type{<:LogDensityFunction{V,M,C,AD}}
-) where {V,M,C,AD<:ADTypes.AbstractADType}
+    ::Type{<:LogDensityFunction{M,V,C,AD}}
+) where {M,V,C,AD<:ADTypes.AbstractADType}
     return LogDensityProblems.LogDensityOrder{1}()
 end
 function LogDensityProblems.logdensity(f::LogDensityFunction, x::AbstractVector)
     return logdensity_at(x, f.model, f.varinfo, f.context)
 end
 function LogDensityProblems.logdensity_and_gradient(
-    f::LogDensityFunction{V,M,C,AD}, x::AbstractVector
-) where {V,M,C,AD<:ADTypes.AbstractADType}
+    f::LogDensityFunction{M,V,C,AD}, x::AbstractVector
+) where {M,V,C,AD<:ADTypes.AbstractADType}
     f.prep === nothing &&
         error("Gradient preparation not available; this should not happen")
     x = map(identity, x)  # Concretise type
@@ -231,11 +242,6 @@ use_closure(::ADTypes.AutoForwardDiff) = false
 use_closure(::ADTypes.AutoMooncake) = false
 use_closure(::ADTypes.AutoReverseDiff) = true
 
-# If a `context` has been specified, we use that. Otherwise we just use the leaf context of `model`.
-function getcontext(f::LogDensityFunction)
-    return f.context === nothing ? leafcontext(f.model.context) : f.context
-end
-
 """
     getmodel(f)
 
@@ -249,7 +255,7 @@ getmodel(f::DynamicPPL.LogDensityFunction) = f.model
 Set the `DynamicPPL.Model` in the given log-density function `f` to `model`.
 """
 function setmodel(f::DynamicPPL.LogDensityFunction, model::DynamicPPL.Model)
-    return Accessors.@set f.model = model
+    return LogDensityFunction(model, f.varinfo, f.context; adtype=f.adtype)
 end
 
 """