Revert "Move src/test_utils and test/test_util to DynamicPPLTestExt"

This reverts commit dcd24e7.

Author: penelopeysm

Project.toml

@@ -31,7 +31,6 @@ EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [extensions]
@@ -40,7 +39,6 @@ DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
 DynamicPPLReverseDiffExt = ["ReverseDiff"]
-DynamicPPLTestExt = ["Test"]
 DynamicPPLZygoteRulesExt = ["ZygoteRules"]
 
 [compat]
@@ -69,3 +67,11 @@ ReverseDiff = "1"
 Test = "1.6"
 ZygoteRules = "0.2"
 julia = "1.10"
+
+[extras]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
+ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"

ext/DynamicPPLTestExt.jl

@@ -1,8 +1,4 @@
-module TestExtUtils
-
-###################################################
-# These used to be in DPPL/src/test_utils.jl ######
-###################################################
+module TestUtils
 
 using AbstractMCMC
 using DynamicPPL
@@ -1101,123 +1097,4 @@ function DynamicPPL.dot_tilde_observe(
     return logp * context.mod, vi
 end
 
-
-
-###################################################
-# These used to be in DPPL/test/test_util.jl ######
-###################################################
-
-# default model
-@model function gdemo_d()
-    s ~ InverseGamma(2, 3)
-    m ~ Normal(0, sqrt(s))
-    1.5 ~ Normal(m, sqrt(s))
-    2.0 ~ Normal(m, sqrt(s))
-    return s, m
-end
-const gdemo_default = gdemo_d()
-
-function test_model_ad(model, logp_manual)
-    vi = VarInfo(model)
-    x = DynamicPPL.getall(vi)
-
-    # Log probabilities using the model.
-    ℓ = DynamicPPL.LogDensityFunction(model, vi)
-    logp_model = Base.Fix1(LogDensityProblems.logdensity, ℓ)
-
-    # Check that both functions return the same values.
-    lp = logp_manual(x)
-    @test logp_model(x) ≈ lp
-
-    # Gradients based on the manual implementation.
-    grad = ForwardDiff.gradient(logp_manual, x)
-
-    y, back = Tracker.forward(logp_manual, x)
-    @test Tracker.data(y) ≈ lp
-    @test Tracker.data(back(1)[1]) ≈ grad
-
-    y, back = Zygote.pullback(logp_manual, x)
-    @test y ≈ lp
-    @test back(1)[1] ≈ grad
-
-    # Gradients based on the model.
-    @test ForwardDiff.gradient(logp_model, x) ≈ grad
-
-    y, back = Tracker.forward(logp_model, x)
-    @test Tracker.data(y) ≈ lp
-    @test Tracker.data(back(1)[1]) ≈ grad
-
-    y, back = Zygote.pullback(logp_model, x)
-    @test y ≈ lp
-    @test back(1)[1] ≈ grad
-end
-
-"""
-    test_setval!(model, chain; sample_idx = 1, chain_idx = 1)
-
-Test `setval!` on `model` and `chain`.
-
-Worth noting that this only supports models containing symbols of the forms
-`m`, `m[1]`, `m[1, 2]`, not `m[1][1]`, etc.
-"""
-function test_setval!(model, chain; sample_idx=1, chain_idx=1)
-    var_info = VarInfo(model)
-    spl = SampleFromPrior()
-    θ_old = var_info[spl]
-    DynamicPPL.setval!(var_info, chain, sample_idx, chain_idx)
-    θ_new = var_info[spl]
-    @test θ_old != θ_new
-    vals = DynamicPPL.values_as(var_info, OrderedDict)
-    iters = map(DynamicPPL.varname_and_value_leaves, keys(vals), values(vals))
-    for (n, v) in mapreduce(collect, vcat, iters)
-        n = string(n)
-        if Symbol(n) ∉ keys(chain)
-            # Assume it's a group
-            chain_val = vec(
-                MCMCChains.group(chain, Symbol(n)).value[sample_idx, :, chain_idx]
-            )
-            v_true = vec(v)
-        else
-            chain_val = chain[sample_idx, n, chain_idx]
-            v_true = v
-        end
-
-        @test v_true == chain_val
-    end
-end
-
-"""
-    short_varinfo_name(vi::AbstractVarInfo)
-
-Return string representing a short description of `vi`.
-"""
-short_varinfo_name(vi::DynamicPPL.ThreadSafeVarInfo) =
-    "threadsafe($(short_varinfo_name(vi.varinfo)))"
-function short_varinfo_name(vi::TypedVarInfo)
-    DynamicPPL.has_varnamedvector(vi) && return "TypedVarInfo with VarNamedVector"
-    return "TypedVarInfo"
-end
-short_varinfo_name(::UntypedVarInfo) = "UntypedVarInfo"
-short_varinfo_name(::DynamicPPL.VectorVarInfo) = "VectorVarInfo"
-short_varinfo_name(::SimpleVarInfo{<:NamedTuple}) = "SimpleVarInfo{<:NamedTuple}"
-short_varinfo_name(::SimpleVarInfo{<:OrderedDict}) = "SimpleVarInfo{<:OrderedDict}"
-function short_varinfo_name(::SimpleVarInfo{<:DynamicPPL.VarNamedVector})
-    return "SimpleVarInfo{<:VarNamedVector}"
-end
-
-# convenient functions for testing model.jl
-# function to modify the representation of values based on their length
-function modify_value_representation(nt::NamedTuple)
-    modified_nt = NamedTuple()
-    for (key, value) in zip(keys(nt), values(nt))
-        if length(value) == 1  # Scalar value
-            modified_value = value[1]
-        else  # Non-scalar value
-            modified_value = value
-        end
-        modified_nt = merge(modified_nt, (key => modified_value,))
-    end
-    return modified_nt
 end
-
-end  # module TestExtUtils

ext/DynamicPPLTestExt/utils.jl renamed to src/test_utils.jl

@@ -1,11 +1,11 @@
 @testset "AD: ForwardDiff and ReverseDiff" begin
-    @testset "$(m.f)" for m in TU.DEMO_MODELS
+    @testset "$(m.f)" for m in DynamicPPL.TestUtils.DEMO_MODELS
         f = DynamicPPL.LogDensityFunction(m)
-        rand_param_values = TU.rand_prior_true(m)
-        vns = TU.varnames(m)
-        varinfos = TU.setup_varinfos(m, rand_param_values, vns)
+        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 "$(TU.short_varinfo_name(varinfo))" for varinfo in varinfos
+        @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
             f = DynamicPPL.LogDensityFunction(m, varinfo)
 
             # use ForwardDiff result as reference

test/ad.jl

@@ -12,7 +12,7 @@
                    logpdf(dist, 2.0)
         end
 
-        TU.test_model_ad(TU.gdemo_default, logp_gdemo_default)
+        test_model_ad(gdemo_default, logp_gdemo_default)
 
         @model function wishart_ad()
             return v ~ Wishart(7, [1 0.5; 0.5 1])
@@ -24,7 +24,7 @@
             return logpdf(dist, reshape(x, 2, 2))
         end
 
-        TU.test_model_ad(wishart_ad(), logp_wishart_ad)
+        test_model_ad(wishart_ad(), logp_wishart_ad)
     end
 
     # https://github.com/TuringLang/Turing.jl/issues/1595

test/compat/ad.jl

@@ -168,7 +168,7 @@ end
 
                     # Let's check elementwise.
                     for vn_child in
-                        TU.varname_leaves(vn_without_prefix, val)
+                        DynamicPPL.TestUtils.varname_leaves(vn_without_prefix, val)
                         if getoptic(vn_child)(val) === missing
                             @test contextual_isassumption(context, vn_child)
                         else
@@ -201,7 +201,7 @@ end
                     vn_without_prefix = remove_prefix(vn)
 
                     for vn_child in
-                        TU.varname_leaves(vn_without_prefix, val)
+                        DynamicPPL.TestUtils.varname_leaves(vn_without_prefix, val)
                         # `vn_child` should be in `context`.
                         @test hasconditioned_nested(context, vn_child)
                         # Value should be the same as extracted above.
@@ -216,7 +216,7 @@ end
     @testset "Evaluation" begin
         @testset "$context" for context in contexts
             # Just making sure that we can actually sample with each of the contexts.
-            @test (TU.gdemo_default(SamplingContext(context)); true)
+            @test (gdemo_default(SamplingContext(context)); true)
         end
     end
 
@@ -258,7 +258,7 @@ end
     end
 
     @testset "FixedContext" begin
-        @testset "$(model.f)" for model in TU.DEMO_MODELS
+        @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
             retval = model()
             s, m = retval.s, retval.m
 

test/contexts.jl

@@ -1,19 +1,19 @@
 @testset "check_model" begin
     @testset "context interface" begin
         # HACK: Require a model to instantiate it, so let's just grab one.
-        model = first(TU.DEMO_MODELS)
+        model = first(DynamicPPL.TestUtils.DEMO_MODELS)
         context = DynamicPPL.DebugUtils.DebugContext(model)
-        TU.test_context_interface(context)
+        DynamicPPL.TestUtils.test_context_interface(context)
     end
 
-    @testset "$(model.f)" for model in TU.DEMO_MODELS
+    @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
         issuccess, trace = check_model_and_trace(model)
         # These models should all work.
         @test issuccess
 
         # Check that the trace contains all the variables in the model.
         varnames_in_trace = DynamicPPL.DebugUtils.varnames_in_trace(trace)
-        for vn in TU.varnames(model)
+        for vn in DynamicPPL.TestUtils.varnames(model)
             @test vn in varnames_in_trace
         end
 
@@ -156,7 +156,7 @@
     end
 
     @testset "comparing multiple traces" begin
-        model = TU.demo_dynamic_constraint()
+        model = DynamicPPL.TestUtils.demo_dynamic_constraint()
         issuccess_1, trace_1 = check_model_and_trace(model)
         issuccess_2, trace_2 = check_model_and_trace(model)
         @test issuccess_1 && issuccess_2

test/debug_utils.jl

@@ -75,8 +75,8 @@ end
         model = demo()
 
         example_values = rand(NamedTuple, model)
-        vis = TU.setup_varinfos(model, example_values, (@varname(m),))
-        @testset "$(TU.short_varinfo_name(vi))" for vi in vis
+        vis = DynamicPPL.TestUtils.setup_varinfos(model, example_values, (@varname(m),))
+        @testset "$(short_varinfo_name(vi))" for vi in vis
             # Evaluate once to ensure we have `logp` value.
             vi = last(DynamicPPL.evaluate!!(model, vi, DefaultContext()))
             vi_linked = if mutable
@@ -109,10 +109,10 @@ end
                 model = demo_lkj(d)
                 dist = LKJCholesky(d, 1.0, uplo)
                 values_original = rand(NamedTuple, model)
-                vis = TU.setup_varinfos(
+                vis = DynamicPPL.TestUtils.setup_varinfos(
                     model, values_original, (@varname(x),)
                 )
-                @testset "$(TU.short_varinfo_name(vi))" for vi in vis
+                @testset "$(short_varinfo_name(vi))" for vi in vis
                     val = vi[@varname(x), dist]
                     # Ensure that `reconstruct` works as intended.
                     @test val isa Cholesky
@@ -150,8 +150,8 @@ end
         @testset "d=$d" for d in [2, 3, 5]
             model = demo_dirichlet(d)
             example_values = rand(NamedTuple, model)
-            vis = TU.setup_varinfos(model, example_values, (@varname(x),))
-            @testset "$(TU.short_varinfo_name(vi))" for vi in vis
+            vis = DynamicPPL.TestUtils.setup_varinfos(model, example_values, (@varname(x),))
+            @testset "$(short_varinfo_name(vi))" for vi in vis
                 lp = logpdf(Dirichlet(d, 1.0), vi[:])
                 @test length(vi[:]) == d
                 lp_model = logjoint(model, vi)
@@ -189,8 +189,8 @@ end
         ]
             model = demo_highdim_dirichlet(ns...)
             example_values = rand(NamedTuple, model)
-            vis = TU.setup_varinfos(model, example_values, (@varname(x),))
-            @testset "$(TU.short_varinfo_name(vi))" for vi in vis
+            vis = DynamicPPL.TestUtils.setup_varinfos(model, example_values, (@varname(x),))
+            @testset "$(short_varinfo_name(vi))" for vi in vis
                 # Linked.
                 vi_linked = if mutable
                     DynamicPPL.link!!(deepcopy(vi), model)

test/linking.jl

@@ -1,8 +1,8 @@
 using Test, DynamicPPL, ADTypes, LogDensityProblems, LogDensityProblemsAD, ReverseDiff
 
 @testset "`getmodel` and `setmodel`" begin
-    @testset "$(nameof(model))" for model in TU.DEMO_MODELS
-        model = TU.DEMO_MODELS[1]
+    @testset "$(nameof(model))" for model in DynamicPPL.TestUtils.DEMO_MODELS
+        model = DynamicPPL.TestUtils.DEMO_MODELS[1]
         ℓ = DynamicPPL.LogDensityFunction(model)
         @test DynamicPPL.getmodel(ℓ) == model
         @test DynamicPPL.setmodel(ℓ, model).model == model
@@ -21,10 +21,10 @@ using Test, DynamicPPL, ADTypes, LogDensityProblems, LogDensityProblemsAD, Rever
 end
 
 @testset "LogDensityFunction" begin
-    @testset "$(nameof(model))" for model in TU.DEMO_MODELS
-        example_values = TU.rand_prior_true(model)
-        vns = TU.varnames(model)
-        varinfos = TU.setup_varinfos(model, example_values, vns)
+    @testset "$(nameof(model))" for model in DynamicPPL.TestUtils.DEMO_MODELS
+        example_values = DynamicPPL.TestUtils.rand_prior_true(model)
+        vns = DynamicPPL.TestUtils.varnames(model)
+        varinfos = DynamicPPL.TestUtils.setup_varinfos(model, example_values, vns)
 
         @testset "$(varinfo)" for varinfo in varinfos
             logdensity = DynamicPPL.LogDensityFunction(model, varinfo)