Fix strictness failure with DifferentiationInterface 0.7

Author: penelopeysm

src/logdensityfunction.jl

@@ -106,6 +106,8 @@ struct LogDensityFunction{
     adtype::AD
     "(internal use only) gradient preparation object for the model"
     prep::Union{Nothing,DI.GradientPrep}
+    "(internal use only) the closure used for the gradient preparation"
+    closure::Union{Nothing,Function}
 
     function LogDensityFunction(
         model::Model,
@@ -114,6 +116,7 @@ struct LogDensityFunction{
         adtype::Union{ADTypes.AbstractADType,Nothing}=nothing,
     )
         if adtype === nothing
+            closure = nothing
             prep = nothing
         else
             # Make backend-specific tweaks to the adtype
@@ -124,10 +127,16 @@ struct LogDensityFunction{
             # Get a set of dummy params to use for prep
             x = map(identity, varinfo[:])
             if use_closure(adtype)
-                prep = DI.prepare_gradient(
-                    x -> logdensity_at(x, model, varinfo, context), adtype, x
-                )
+                # The closure itself has to be stored inside the
+                # LogDensityFunction to ensure that the signature of the
+                # function being differentiated is the same as that used for
+                # preparation. See
+                # https://github.com/TuringLang/DynamicPPL.jl/pull/922 for an
+                # explanation.
+                closure = x -> logdensity_at(x, model, varinfo, context)
+                prep = DI.prepare_gradient(closure, adtype, x)
             else
+                closure = nothing
                 prep = DI.prepare_gradient(
                     logdensity_at,
                     adtype,
@@ -139,7 +148,7 @@ struct LogDensityFunction{
             end
         end
         return new{typeof(model),typeof(varinfo),typeof(context),typeof(adtype)}(
-            model, varinfo, context, adtype, prep
+            model, varinfo, context, adtype, prep, closure
         )
     end
 end
@@ -208,9 +217,8 @@ function LogDensityProblems.logdensity_and_gradient(
     # Make branching statically inferrable, i.e. type-stable (even if the two
     # branches happen to return different types)
     return if use_closure(f.adtype)
-        DI.value_and_gradient(
-            x -> logdensity_at(x, f.model, f.varinfo, f.context), f.prep, f.adtype, x
-        )
+        f.closure === nothing && error("Closure not available; this should not happen")
+        DI.value_and_gradient(f.closure, f.prep, f.adtype, x)
     else
         DI.value_and_gradient(
             logdensity_at,