Merge remote-tracking branch 'origin/main' into py/no-mooncake-pre

Author: penelopeysm

HISTORY.md

@@ -2,8 +2,9 @@
 
 ## 0.36.4
 
-Added a compatibility entry for JET.jl 0.10.
-This should only affect you if you are using DynamicPPL on the Julia 1.12 pre-release.
+Added compatibility with DifferentiationInterface.jl 0.7, and also with JET.jl 0.10.
+
+The JET compatibility entry should only affect you if you are using DynamicPPL on the Julia 1.12 pre-release.
 
 ## 0.36.3
 

Project.toml

@@ -54,7 +54,7 @@ ChainRulesCore = "1"
 Chairmarks = "1.3.1"
 Compat = "4"
 ConstructionBase = "1.5.4"
-DifferentiationInterface = "0.6.41"
+DifferentiationInterface = "0.6.41, 0.7"
 Distributions = "0.25"
 DocStringExtensions = "0.9"
 EnzymeCore = "0.6 - 0.8"

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -23,7 +23,7 @@ Return the dimension of `model`, accounting for linking, if any.
 """
 function model_dimension(model, islinked)
     vi = VarInfo()
-    model(vi)
+    model(StableRNG(23), vi)
     if islinked
         vi = DynamicPPL.link(vi, model)
     end

src/logdensityfunction.jl

@@ -124,9 +124,7 @@ 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
-                )
+                prep = DI.prepare_gradient(LogDensityAt(model, varinfo, context), adtype, x)
             else
                 prep = DI.prepare_gradient(
                     logdensity_at,
@@ -184,6 +182,26 @@ function logdensity_at(
     return getlogp(last(evaluate!!(model, varinfo_new, context)))
 end
 
+"""
+    LogDensityAt{M<:Model,V<:AbstractVarInfo,C<:AbstractContext}(
+        model::M
+        varinfo::V
+        context::C
+    )
+
+A callable struct that serves the same purpose as `x -> logdensity_at(x, model,
+varinfo, context)`.
+"""
+struct LogDensityAt{M<:Model,V<:AbstractVarInfo,C<:AbstractContext}
+    model::M
+    varinfo::V
+    context::C
+end
+function (ld::LogDensityAt)(x::AbstractVector)
+    varinfo_new = unflatten(ld.varinfo, x)
+    return getlogp(last(evaluate!!(ld.model, varinfo_new, ld.context)))
+end
+
 ### LogDensityProblems interface
 
 function LogDensityProblems.capabilities(
@@ -209,7 +227,7 @@ function LogDensityProblems.logdensity_and_gradient(
     # 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
+            LogDensityAt(f.model, f.varinfo, f.context), f.prep, f.adtype, x
         )
     else
         DI.value_and_gradient(

test/Project.toml

@@ -36,7 +36,7 @@ Aqua = "0.8"
 Bijectors = "0.15.1"
 Combinatorics = "1"
 Compat = "4.3.0"
-DifferentiationInterface = "0.6.41"
+DifferentiationInterface = "0.6.41, 0.7"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 EnzymeCore = "0.6 - 0.8"