[no ci] The Rest

Author: penelopeysm

docs/src/api.md

@@ -461,7 +461,6 @@ The behaviour of a model execution can be changed with evaluation contexts, whic
 Contexts are subtypes of `AbstractPPL.AbstractContext`.
 
 ```@docs
-SamplingContext
 DefaultContext
 PrefixContext
 ConditionContext
@@ -490,15 +489,7 @@ DynamicPPL.init
 
 ### Samplers
 
-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.
-
-```@docs
-SampleFromPrior
-SampleFromUniform
-```
-
-Additionally, a generic sampler for inference is implemented.
+In DynamicPPL a generic sampler for inference is implemented.
 
 ```@docs
 Sampler
@@ -509,7 +500,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
 

src/DynamicPPL.jl

@@ -97,13 +97,10 @@ export AbstractVarInfo,
     values_as_in_model,
     # Samplers
     Sampler,
-    SampleFromPrior,
-    SampleFromUniform,
     # LogDensityFunction
     LogDensityFunction,
     # Contexts
     contextualize,
-    SamplingContext,
     DefaultContext,
     PrefixContext,
     ConditionContext,

src/context_implementations.jl

@@ -1,38 +1,11 @@
 # assume
-"""
-    tilde_assume(context::SamplingContext, right, vn, vi)
-
-Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
-accumulate the log probability, and return the sampled value with a context associated
-with a sampler.
-
-Falls back to
-```julia
-tilde_assume(context.rng, context.context, context.sampler, right, vn, vi)
-```
-"""
-function tilde_assume(context::SamplingContext, right, vn, vi)
-    return tilde_assume(context.rng, context.context, context.sampler, right, vn, vi)
-end
-
 function tilde_assume(context::AbstractContext, args...)
     return tilde_assume(childcontext(context), args...)
 end
 function tilde_assume(::DefaultContext, right, vn, vi)
     return assume(right, vn, vi)
 end
 
-function tilde_assume(rng::Random.AbstractRNG, context::AbstractContext, args...)
-    return tilde_assume(rng, childcontext(context), args...)
-end
-function tilde_assume(rng::Random.AbstractRNG, ::DefaultContext, sampler, right, vn, vi)
-    return assume(rng, sampler, right, vn, vi)
-end
-function tilde_assume(::DefaultContext, sampler, right, vn, vi)
-    # same as above but no rng
-    return assume(Random.default_rng(), sampler, right, vn, vi)
-end
-
 function tilde_assume(context::PrefixContext, right, vn, vi)
     # Note that we can't use something like this here:
     #     new_vn = prefix(context, vn)
@@ -46,12 +19,6 @@ function tilde_assume(context::PrefixContext, right, vn, vi)
     new_vn, new_context = prefix_and_strip_contexts(context, vn)
     return tilde_assume(new_context, right, new_vn, vi)
 end
-function tilde_assume(
-    rng::Random.AbstractRNG, context::PrefixContext, sampler, right, vn, vi
-)
-    new_vn, new_context = prefix_and_strip_contexts(context, vn)
-    return tilde_assume(rng, new_context, sampler, right, new_vn, vi)
-end
 
 """
     tilde_assume!!(context, right, vn, vi)
@@ -71,17 +38,6 @@ function tilde_assume!!(context, right, vn, vi)
 end
 
 # observe
-"""
-    tilde_observe!!(context::SamplingContext, right, left, vi)
-
-Handle observed constants with a `context` associated with a sampler.
-
-Falls back to `tilde_observe!!(context.context, right, left, vi)`.
-"""
-function tilde_observe!!(context::SamplingContext, right, left, vn, vi)
-    return tilde_observe!!(context.context, right, left, vn, vi)
-end
-
 function tilde_observe!!(context::AbstractContext, right, left, vn, vi)
     return tilde_observe!!(childcontext(context), right, left, vn, vi)
 end
@@ -115,10 +71,6 @@ function tilde_observe!!(::DefaultContext, right, left, vn, vi)
     return left, vi
 end
 
-function assume(::Random.AbstractRNG, spl::Sampler, dist)
-    return error("DynamicPPL.assume: unmanaged inference algorithm: $(typeof(spl))")
-end
-
 # fallback without sampler
 function assume(dist::Distribution, vn::VarName, vi)
     y = getindex_internal(vi, vn)
@@ -127,46 +79,3 @@ function assume(dist::Distribution, vn::VarName, vi)
     vi = accumulate_assume!!(vi, x, logjac, vn, dist)
     return x, vi
 end
-
-# TODO: Remove this thing.
-# SampleFromPrior and SampleFromUniform
-function assume(
-    rng::Random.AbstractRNG,
-    sampler::Union{SampleFromPrior,SampleFromUniform},
-    dist::Distribution,
-    vn::VarName,
-    vi::VarInfoOrThreadSafeVarInfo,
-)
-    if haskey(vi, vn)
-        # Always overwrite the parameters with new ones for `SampleFromUniform`.
-        if sampler isa SampleFromUniform || is_flagged(vi, vn, "del")
-            # TODO(mhauru) Is it important to unset the flag here? The `true` allows us
-            # to ignore the fact that for VarNamedVector this does nothing, but I'm unsure
-            # if that's okay.
-            unset_flag!(vi, vn, "del", true)
-            r = init(rng, dist, sampler)
-            f = to_maybe_linked_internal_transform(vi, vn, dist)
-            # TODO(mhauru) This should probably be call a function called setindex_internal!
-            vi = BangBang.setindex!!(vi, f(r), vn)
-            setorder!(vi, vn, get_num_produce(vi))
-        else
-            # Otherwise we just extract it.
-            r = vi[vn, dist]
-        end
-    else
-        r = init(rng, dist, sampler)
-        if istrans(vi)
-            f = to_linked_internal_transform(vi, vn, dist)
-            vi = push!!(vi, vn, f(r), dist)
-            # By default `push!!` sets the transformed flag to `false`.
-            vi = settrans!!(vi, true, vn)
-        else
-            vi = push!!(vi, vn, r, dist)
-        end
-    end
-
-    # HACK: The above code might involve an `invlink` somewhere, etc. so we need to correct.
-    logjac = logabsdetjac(istrans(vi, vn) ? link_transform(dist) : identity, r)
-    vi = accumulate_assume!!(vi, r, -logjac, vn, dist)
-    return r, vi
-end

src/contexts.jl

@@ -47,7 +47,7 @@ effectively updating the child context.
 ```jldoctest
 julia> using DynamicPPL: DynamicTransformationContext
 
-julia> ctx = SamplingContext();
+julia> ctx = ConditionContext((; a = 1);
 
 julia> DynamicPPL.childcontext(ctx)
 DefaultContext()
@@ -121,73 +121,6 @@ setleafcontext(::IsLeaf, ::IsParent, left, right) = right
 setleafcontext(::IsLeaf, ::IsLeaf, left, right) = right
 
 # Contexts
-"""
-    SamplingContext(
-            [rng::Random.AbstractRNG=Random.default_rng()],
-            [sampler::AbstractSampler=SampleFromPrior()],
-            [context::AbstractContext=DefaultContext()],
-    )
-
-Create a context that allows you to sample parameters with the `sampler` when running the model.
-The `context` determines how the returned log density is computed when running the model.
-
-See also: [`DefaultContext`](@ref)
-"""
-struct SamplingContext{S<:AbstractSampler,C<:AbstractContext,R} <: AbstractContext
-    rng::R
-    sampler::S
-    context::C
-end
-
-function SamplingContext(
-    rng::Random.AbstractRNG=Random.default_rng(), sampler::AbstractSampler=SampleFromPrior()
-)
-    return SamplingContext(rng, sampler, DefaultContext())
-end
-
-function SamplingContext(
-    sampler::AbstractSampler, context::AbstractContext=DefaultContext()
-)
-    return SamplingContext(Random.default_rng(), sampler, context)
-end
-
-function SamplingContext(rng::Random.AbstractRNG, context::AbstractContext)
-    return SamplingContext(rng, SampleFromPrior(), context)
-end
-
-function SamplingContext(context::AbstractContext)
-    return SamplingContext(Random.default_rng(), SampleFromPrior(), context)
-end
-
-NodeTrait(context::SamplingContext) = IsParent()
-childcontext(context::SamplingContext) = context.context
-function setchildcontext(parent::SamplingContext, child)
-    return SamplingContext(parent.rng, parent.sampler, child)
-end
-
-"""
-    hassampler(context)
-
-Return `true` if `context` has a sampler.
-"""
-hassampler(::SamplingContext) = true
-hassampler(context::AbstractContext) = hassampler(NodeTrait(context), context)
-hassampler(::IsLeaf, context::AbstractContext) = false
-hassampler(::IsParent, context::AbstractContext) = hassampler(childcontext(context))
-
-"""
-    getsampler(context)
-
-Return the sampler of the context `context`.
-
-This will traverse the context tree until it reaches the first [`SamplingContext`](@ref),
-at which point it will return the sampler of that context.
-"""
-getsampler(context::SamplingContext) = context.sampler
-getsampler(context::AbstractContext) = getsampler(NodeTrait(context), context)
-getsampler(::IsParent, context::AbstractContext) = getsampler(childcontext(context))
-getsampler(::IsLeaf, ::AbstractContext) = error("No sampler found in context")
-
 """
     struct DefaultContext <: AbstractContext end
 

src/debug_utils.jl

@@ -438,9 +438,10 @@ function check_model_and_trace(
     kwargs...,
 )
     # Execute the model with the debug context.
-    debug_context = DebugContext(
-        SamplingContext(rng, model.context); error_on_failure=error_on_failure, kwargs...
+    new_context = DynamicPPL.setleafcontext(
+        model.context, DynamicPPL.InitContext(rng, DynamicPPL.PriorInit())
     )
+    debug_context = DebugContext(new_context; error_on_failure=error_on_failure, kwargs...)
     debug_model = DynamicPPL.contextualize(model, debug_context)
 
     # Perform checks before evaluating the model.