Move FastLDF to experimental for now

Author: penelopeysm

benchmarks/benchmarks.jl

@@ -59,14 +59,14 @@ chosen_combinations = [
         false,
     ),
     ("Smorgasbord", smorgasbord_instance, :typed, :forwarddiff, false),
-    # ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true),
-    # ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true),
-    # ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
-    # ("Smorgasbord", smorgasbord_instance, :typed_vector, :forwarddiff, true),
-    # ("Smorgasbord", smorgasbord_instance, :untyped_vector, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :simple_namedtuple, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :untyped, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :simple_dict, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :typed_vector, :forwarddiff, true),
+    ("Smorgasbord", smorgasbord_instance, :untyped_vector, :forwarddiff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :reversediff, true),
     ("Smorgasbord", smorgasbord_instance, :typed, :mooncake, true),
-    # ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true),
+    ("Smorgasbord", smorgasbord_instance, :typed, :enzyme, true),
     ("Loop univariate 1k", loop_univariate1k, :typed, :mooncake, true),
     ("Multivariate 1k", multivariate1k, :typed, :mooncake, true),
     ("Loop univariate 10k", loop_univariate10k, :typed, :mooncake, true),

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -94,7 +94,9 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
         vi = DynamicPPL.link(vi, model)
     end
 
-    f = DynamicPPL.FastLDF(model, DynamicPPL.getlogjoint_internal, vi; adtype=adbackend)
+    f = DynamicPPL.LogDensityFunction(
+        model, DynamicPPL.getlogjoint_internal, vi; adtype=adbackend
+    )
     # The parameters at which we evaluate f.
     θ = vi[:]
 

src/DynamicPPL.jl

@@ -191,7 +191,6 @@ include("simple_varinfo.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
 include("logdensityfunction.jl")
-include("fastldf.jl")
 include("model_utils.jl")
 include("extract_priors.jl")
 include("values_as_in_model.jl")

src/experimental.jl

@@ -2,6 +2,8 @@ module Experimental
 
 using DynamicPPL: DynamicPPL
 
+include("fastldf.jl")
+
 # This file only defines the names of the functions, and their docstrings. The actual implementations are in `ext/DynamicPPLJETExt.jl`, since we don't want to depend on JET.jl other than as a weak dependency.
 """
     is_suitable_varinfo(model::Model, varinfo::AbstractVarInfo; kwargs...)

src/fastldf.jl renamed to src/fasteval.jl

@@ -4,7 +4,8 @@ using ADTypes: AbstractADType, AutoForwardDiff
 using Chairmarks: @be
 import DifferentiationInterface as DI
 using DocStringExtensions
-using DynamicPPL: Model, FastLDF, VarInfo, AbstractVarInfo, getlogjoint_internal, link
+using DynamicPPL:
+    Model, LogDensityFunction, VarInfo, AbstractVarInfo, getlogjoint_internal, link
 using LogDensityProblems: logdensity, logdensity_and_gradient
 using Random: AbstractRNG, default_rng
 using Statistics: median
@@ -264,7 +265,7 @@ function run_ad(
     # Calculate log-density and gradient with the backend of interest
     verbose && @info "Running AD on $(model.f) with $(adtype)\n"
     verbose && println("       params : $(params)")
-    ldf = FastLDF(model, getlogdensity, varinfo; adtype=adtype)
+    ldf = LogDensityFunction(model, getlogdensity, varinfo; adtype=adtype)
 
     value, grad = logdensity_and_gradient(ldf, params)
     # collect(): https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754
@@ -281,7 +282,9 @@ function run_ad(
             value_true = test.value
             grad_true = test.grad
         elseif test isa WithBackend
-            ldf_reference = FastLDF(model, getlogdensity, varinfo; adtype=test.adtype)
+            ldf_reference = LogDensityFunction(
+                model, getlogdensity, varinfo; adtype=test.adtype
+            )
             value_true, grad_true = logdensity_and_gradient(ldf_reference, params)
             # collect(): https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754
             grad_true = collect(grad_true)

src/test_utils/ad.jl

@@ -1,6 +1,5 @@
-using DynamicPPL: FastLDF
+using DynamicPPL: LogDensityFunction
 using DynamicPPL.TestUtils.AD: run_ad, WithExpectedResult, NoTest
-using Random: Xoshiro
 
 @testset "Automatic differentiation" begin
     # Used as the ground truth that others are compared against.
@@ -16,25 +15,64 @@ using Random: Xoshiro
         [AutoReverseDiff(; compile=false), AutoReverseDiff(; compile=true)]
     end
 
+    @testset "Unsupported backends" begin
+        @model demo() = x ~ Normal()
+        @test_logs (:warn, r"not officially supported") LogDensityFunction(
+            demo(); adtype=AutoZygote()
+        )
+    end
+
     @testset "Correctness" begin
         @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
-            varinfo = VarInfo(Xoshiro(468), m)
-            linked_varinfo = DynamicPPL.link(varinfo, m)
-            f = FastLDF(m, getlogjoint_internal, linked_varinfo)
-            x = linked_varinfo[:]
-
-            # Calculate reference logp + gradient of logp using ForwardDiff
-            ref_ad_result = run_ad(m, ref_adtype; varinfo=linked_varinfo, test=NoTest())
-            ref_logp, ref_grad = ref_ad_result.value_actual, ref_ad_result.grad_actual
-
-            @testset "$adtype" for adtype in test_adtypes
-                @info "Testing AD on: $(m.f) - $adtype"
-                @test run_ad(
-                    m,
-                    adtype;
-                    varinfo=linked_varinfo,
-                    test=WithExpectedResult(ref_logp, ref_grad),
-                ) isa Any
+            rand_param_values = DynamicPPL.TestUtils.rand_prior_true(m)
+            vns = DynamicPPL.TestUtils.varnames(m)
+            varinfos = DynamicPPL.TestUtils.setup_varinfos(m, rand_param_values, vns)
+
+            @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
+                linked_varinfo = DynamicPPL.link(varinfo, m)
+                f = LogDensityFunction(m, getlogjoint_internal, linked_varinfo)
+                x = DynamicPPL.getparams(f)
+
+                # Calculate reference logp + gradient of logp using ForwardDiff
+                ref_ad_result = run_ad(m, ref_adtype; varinfo=linked_varinfo, test=NoTest())
+                ref_logp, ref_grad = ref_ad_result.value_actual, ref_ad_result.grad_actual
+
+                @testset "$adtype" for adtype in test_adtypes
+                    @info "Testing AD on: $(m.f) - $(short_varinfo_name(linked_varinfo)) - $adtype"
+
+                    # Put predicates here to avoid long lines
+                    is_mooncake = adtype isa AutoMooncake
+                    is_1_10 = v"1.10" <= VERSION < v"1.11"
+                    is_1_11 = v"1.11" <= VERSION < v"1.12"
+                    is_svi_vnv =
+                        linked_varinfo isa SimpleVarInfo{<:DynamicPPL.VarNamedVector}
+                    is_svi_od = linked_varinfo isa SimpleVarInfo{<:OrderedDict}
+
+                    # 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.LogDensityFunction(
+                            m, getlogjoint_internal, linked_varinfo; adtype=adtype
+                        )
+                    elseif is_mooncake && is_1_10 && is_svi_vnv
+                        # TODO: report upstream
+                        @test_throws UndefRefError DynamicPPL.LogDensityFunction(
+                            m, getlogjoint_internal, linked_varinfo; adtype=adtype
+                        )
+                    elseif is_mooncake && is_1_10 && is_svi_od
+                        # TODO: report upstream
+                        @test_throws Mooncake.MooncakeRuleCompilationError DynamicPPL.LogDensityFunction(
+                            m, getlogjoint_internal, linked_varinfo; adtype=adtype
+                        )
+                    else
+                        @test run_ad(
+                            m,
+                            adtype;
+                            varinfo=linked_varinfo,
+                            test=WithExpectedResult(ref_logp, ref_grad),
+                        ) isa Any
+                    end
+                end
             end
         end
     end
@@ -45,7 +83,7 @@ using Random: Xoshiro
         test_m = randn(2, 3)
 
         function eval_logp_and_grad(model, m, adtype)
-            ldf = FastLDF(model(); adtype=adtype)
+            ldf = LogDensityFunction(model(); adtype=adtype)
             return LogDensityProblems.logdensity_and_gradient(ldf, m[:])
         end