Excise SamplingContext tests

Author: penelopeysm

docs/src/api.md

@@ -456,33 +456,24 @@ AbstractPPL.evaluate!!
 
 This method mutates the `varinfo` used for execution.
 By default, it does not perform any actual sampling: it only evaluates the model using the values of the variables that are already in the `varinfo`.
-To perform sampling, you can either wrap `model.context` in a `SamplingContext`, or use this convenience method:
-
-```@docs
-DynamicPPL.evaluate_and_sample!!
-```
 
 The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
 Contexts are subtypes of `AbstractPPL.AbstractContext`.
 
 ```@docs
-SamplingContext
 DefaultContext
 PrefixContext
 ConditionContext
+InitContext
 ```
 
-### Samplers
+### VarInfo initialisation
 
-In DynamicPPL two samplers are defined that are used to initialize unobserved random variables:
-[`SampleFromPrior`](@ref) which samples from the prior distribution, and [`SampleFromUniform`](@ref) which samples from a uniform distribution.
+TODO
 
-```@docs
-SampleFromPrior
-SampleFromUniform
-```
+### Samplers
 
-Additionally, a generic sampler for inference is implemented.
+In DynamicPPL a generic sampler for inference is implemented.
 
 ```@docs
 Sampler
@@ -493,7 +484,7 @@ The default implementation of [`Sampler`](@ref) uses the following unexported fu
 ```@docs
 DynamicPPL.initialstep
 DynamicPPL.loadstate
-DynamicPPL.initialsampler
+DynamicPPL.init_strategy
 ```
 
 Finally, to specify which varinfo type a [`Sampler`](@ref) should use for a given [`Model`](@ref), this is specified by [`DynamicPPL.default_varinfo`](@ref) and can thus be overloaded for each  `model`-`sampler` combination. This can be useful in cases where one has explicit knowledge that one type of varinfo will be more performant for the given `model` and `sampler`.

ext/DynamicPPLEnzymeCoreExt.jl

@@ -8,8 +8,6 @@ else
     using ..EnzymeCore
 end
 
-@inline EnzymeCore.EnzymeRules.inactive_type(::Type{<:DynamicPPL.SamplingContext}) = true
-
 # Mark istrans as having 0 derivative. The `nothing` return value is not significant, Enzyme
 # only checks whether such a method exists, and never runs it.
 @inline EnzymeCore.EnzymeRules.inactive_noinl(::typeof(DynamicPPL.istrans), args...) =

ext/DynamicPPLJETExt.jl

@@ -21,17 +21,12 @@ end
 function DynamicPPL.Experimental._determine_varinfo_jet(
     model::DynamicPPL.Model; only_ddpl::Bool=true
 )
-    # Use SamplingContext to test type stability.
-    sampling_model = DynamicPPL.contextualize(
-        model, DynamicPPL.SamplingContext(model.context)
-    )
-
     # First we try with the typed varinfo.
-    varinfo = DynamicPPL.typed_varinfo(sampling_model)
+    varinfo = DynamicPPL.typed_varinfo(model)
 
     # Let's make sure that both evaluation and sampling doesn't result in type errors.
     issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
-        sampling_model, varinfo; only_ddpl
+        model, varinfo; only_ddpl
     )
 
     if !issuccess
@@ -46,7 +41,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
     else
         # Warn the user that we can't use the type stable one.
         @warn "Model seems incompatible with typed varinfo. Falling back to untyped varinfo."
-        DynamicPPL.untyped_varinfo(sampling_model)
+        DynamicPPL.untyped_varinfo(model)
     end
 end
 

test/ad.jl

@@ -112,8 +112,7 @@ using DynamicPPL.TestUtils.AD: run_ad, WithExpectedResult, NoTest
         # Compiling the ReverseDiff tape used to fail here
         spl = Sampler(MyEmptyAlg())
         vi = VarInfo(model)
-        sampling_model = contextualize(model, SamplingContext(model.context))
-        ldf = LogDensityFunction(sampling_model, vi; adtype=AutoReverseDiff(; compile=true))
+        ldf = LogDensityFunction(model, vi; adtype=AutoReverseDiff(; compile=true))
         @test LogDensityProblems.logdensity_and_gradient(ldf, vi[:]) isa Any
     end
 

test/compiler.jl

@@ -193,11 +193,8 @@ module Issue537 end
         varinfo = VarInfo(model)
         @test getlogjoint(varinfo) == lp
         @test varinfo_ isa AbstractVarInfo
-        # During the model evaluation, its context is wrapped in a
-        # SamplingContext, so `model_` is not going to be equal to `model`.
-        # We can still check equality of `f` though.
         @test model_.f === model.f
-        @test model_.context isa SamplingContext
+        @test model_.context isa InitContext
         @test model_.context.rng isa Random.AbstractRNG
 
         # disable warnings

test/contexts.jl

@@ -49,7 +49,6 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
     contexts = Dict(
         :default => DefaultContext(),
         :testparent => DynamicPPL.TestUtils.TestParentContext(DefaultContext()),
-        :sampling => SamplingContext(),
         :prefix => PrefixContext(@varname(x)),
         :condition1 => ConditionContext((x=1.0,)),
         :condition2 => ConditionContext(
@@ -150,11 +149,11 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
             vn = @varname(x[1])
             ctx1 = PrefixContext(@varname(a))
             @test DynamicPPL.prefix(ctx1, vn) == @varname(a.x[1])
-            ctx2 = SamplingContext(ctx1)
+            ctx2 = ConditionContext(Dict(), ctx1)
             @test DynamicPPL.prefix(ctx2, vn) == @varname(a.x[1])
             ctx3 = PrefixContext(@varname(b), ctx2)
             @test DynamicPPL.prefix(ctx3, vn) == @varname(b.a.x[1])
-            ctx4 = DynamicPPL.SamplingContext(ctx3)
+            ctx4 = FixedContext(Dict(), ctx3)
             @test DynamicPPL.prefix(ctx4, vn) == @varname(b.a.x[1])
         end
 
@@ -165,29 +164,28 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
             @test new_vn == @varname(a.x[1])
             @test new_ctx == DefaultContext()
 
-            ctx2 = SamplingContext(PrefixContext(@varname(a)))
+            ctx2 = FixedContext((b=4,), PrefixContext(@varname(a)))
             new_vn, new_ctx = DynamicPPL.prefix_and_strip_contexts(ctx2, vn)
             @test new_vn == @varname(a.x[1])
-            @test new_ctx == SamplingContext()
+            @test new_ctx == FixedContext((b=4,))
 
             ctx3 = PrefixContext(@varname(a), ConditionContext((a=1,)))
             new_vn, new_ctx = DynamicPPL.prefix_and_strip_contexts(ctx3, vn)
             @test new_vn == @varname(a.x[1])
             @test new_ctx == ConditionContext((a=1,))
 
-            ctx4 = SamplingContext(PrefixContext(@varname(a), ConditionContext((a=1,))))
+            ctx4 = FixedContext((b=4,)PrefixContext(@varname(a), ConditionContext((a=1,))))
             new_vn, new_ctx = DynamicPPL.prefix_and_strip_contexts(ctx4, vn)
             @test new_vn == @varname(a.x[1])
-            @test new_ctx == SamplingContext(ConditionContext((a=1,)))
+            @test new_ctx == FixedContext((b=4,)ConditionContext((a=1,)))
         end
 
         @testset "evaluation: $(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
             prefix_vn = @varname(my_prefix)
-            context = DynamicPPL.PrefixContext(prefix_vn, SamplingContext())
-            sampling_model = contextualize(model, context)
-            # Sample with the context.
-            varinfo = DynamicPPL.VarInfo()
-            DynamicPPL.evaluate!!(sampling_model, varinfo)
+            context = DynamicPPL.PrefixContext(prefix_vn, DefaultContext())
+            new_model = contextualize(model, context)
+            # Initialize a new varinfo with the prefixed model
+            DynamicPPL.init!!(new_model, DynamicPPL.VarInfo())
             # Extract the resulting varnames
             vns_actual = Set(keys(varinfo))
 
@@ -202,22 +200,6 @@ Base.IteratorEltype(::Type{<:AbstractContext}) = Base.EltypeUnknown()
         end
     end
 
-    @testset "SamplingContext" begin
-        context = SamplingContext(Random.default_rng(), SampleFromPrior(), DefaultContext())
-        @test context isa SamplingContext
-
-        # convenience constructors
-        @test SamplingContext() == context
-        @test SamplingContext(Random.default_rng()) == context
-        @test SamplingContext(SampleFromPrior()) == context
-        @test SamplingContext(DefaultContext()) == context
-        @test SamplingContext(Random.default_rng(), SampleFromPrior()) == context
-        @test SamplingContext(Random.default_rng(), DefaultContext()) == context
-        @test SamplingContext(SampleFromPrior(), DefaultContext()) == context
-        @test SamplingContext(SampleFromPrior(), DefaultContext()) == context
-        @test EnzymeCore.EnzymeRules.inactive_type(typeof(context))
-    end
-
     @testset "ConditionContext" begin
         @testset "Nesting" begin
             @testset "NamedTuple" begin

test/ext/DynamicPPLJETExt.jl

@@ -62,17 +62,16 @@
 
     @testset "demo models" begin
         @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
-            sampling_model = contextualize(model, SamplingContext(model.context))
             # Use debug logging below.
             varinfo = DynamicPPL.Experimental.determine_suitable_varinfo(model)
-            # Check that the inferred varinfo is indeed suitable for evaluation and sampling
+            # Check that the inferred varinfo is indeed suitable for evaluation and initialisation
             f_eval, argtypes_eval = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
                 model, varinfo
             )
             JET.test_call(f_eval, argtypes_eval)
 
             f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                sampling_model, varinfo
+                init_model, varinfo
             )
             JET.test_call(f_sample, argtypes_sample)
             # For our demo models, they should all result in typed.
@@ -85,10 +84,6 @@
                     model, typed_vi
                 )
                 JET.test_call(f_eval, argtypes_eval)
-                f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                    sampling_model, typed_vi
-                )
-                JET.test_call(f_sample, argtypes_sample)
             end
         end
     end

test/threadsafe.jl

@@ -68,8 +68,7 @@
         @time model(vi)
 
         # Ensure that we use `ThreadSafeVarInfo` to handle multithreaded observe statements.
-        sampling_model = contextualize(model, SamplingContext(model.context))
-        DynamicPPL.evaluate_threadsafe!!(sampling_model, vi)
+        DynamicPPL.evaluate_threadsafe!!(model, vi)
         @test getlogjoint(vi) ≈ lp_w_threads
         # check that it's wrapped during the model evaluation
         @test vi_ isa DynamicPPL.ThreadSafeVarInfo
@@ -104,8 +103,7 @@
         @test lp_w_threads ≈ lp_wo_threads
 
         # Ensure that we use `VarInfo`.
-        sampling_model = contextualize(model, SamplingContext(model.context))
-        DynamicPPL.evaluate_threadunsafe!!(sampling_model, vi)
+        DynamicPPL.evaluate_threadunsafe!!(model, vi)
         @test getlogjoint(vi) ≈ lp_w_threads
         @test vi_ isa VarInfo
         @test vi isa VarInfo