[no ci] Remove SamplingContext for good

Author: penelopeysm

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...) =

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,43 +1,14 @@
 # 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(::Random.AbstractRNG, ::InitContext, sampler, right, vn, vi)
-    return error(
-        "Encountered SamplingContext->InitContext. This method will be removed in the next PR.",
-    )
-end
-function tilde_assume(::DefaultContext, sampler, right, vn, vi)
-    # same as above but no rng
-    return assume(Random.default_rng(), sampler, right, vn, vi)
+    y = getindex_internal(vi, vn)
+    f = from_maybe_linked_internal_transform(vi, vn, right)
+    x, logjac = with_logabsdet_jacobian(f, y)
+    vi = accumulate_assume!!(vi, x, logjac, vn, right)
+    return x, 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)
@@ -51,12 +22,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)
@@ -76,17 +41,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
@@ -119,59 +73,3 @@ function tilde_observe!!(::DefaultContext, right, left, vn, vi)
     vi = accumulate_observe!!(vi, right, left, vn)
     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)
-    f = from_maybe_linked_internal_transform(vi, vn, dist)
-    x, logjac = with_logabsdet_jacobian(f, y)
-    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.

src/sampler.jl

@@ -1,34 +1,3 @@
-# TODO: Make `UniformSampling` and `Prior` algs + just use `Sampler`
-# That would let us use all defaults for Sampler, combine it with other samplers etc.
-"""
-    SampleFromUniform
-
-Sampling algorithm that samples unobserved random variables from a uniform distribution.
-
-# References
-
-[Stan reference manual](https://mc-stan.org/docs/2_28/reference-manual/initialization.html#random-initial-values)
-"""
-struct SampleFromUniform <: AbstractSampler end
-
-"""
-    SampleFromPrior
-
-Sampling algorithm that samples unobserved random variables from their prior distribution.
-"""
-struct SampleFromPrior <: AbstractSampler end
-
-# Initializations.
-init(rng, dist, ::SampleFromPrior) = rand(rng, dist)
-function init(rng, dist, ::SampleFromUniform)
-    return istransformable(dist) ? inittrans(rng, dist) : rand(rng, dist)
-end
-
-init(rng, dist, ::SampleFromPrior, n::Int) = rand(rng, dist, n)
-function init(rng, dist, ::SampleFromUniform, n::Int)
-    return istransformable(dist) ? inittrans(rng, dist, n) : rand(rng, dist, n)
-end
-
 # TODO(mhauru) Could we get rid of Sampler now that it's just a wrapper around `alg`?
 # (Selector has been removed).
 """
@@ -49,20 +18,6 @@ struct Sampler{T} <: AbstractSampler
     alg::T
 end
 
-# AbstractMCMC interface for SampleFromUniform and SampleFromPrior
-function AbstractMCMC.step(
-    rng::Random.AbstractRNG,
-    model::Model,
-    sampler::Union{SampleFromUniform,SampleFromPrior},
-    state=nothing;
-    kwargs...,
-)
-    vi = VarInfo()
-    strategy = sampler isa SampleFromPrior ? PriorInit() : UniformInit()
-    DynamicPPL.init!!(rng, model, vi, strategy)
-    return vi, nothing
-end
-
 """
     default_varinfo(rng, model, sampler)
 

src/simple_varinfo.jl

@@ -457,23 +457,6 @@ function Base.merge(varinfo_left::SimpleVarInfo, varinfo_right::SimpleVarInfo)
 end
 
 # Context implementations
-# NOTE: Evaluations, i.e. those without `rng` are shared with other
-# implementations of `AbstractVarInfo`.
-function assume(
-    rng::Random.AbstractRNG,
-    sampler::Union{SampleFromPrior,SampleFromUniform},
-    dist::Distribution,
-    vn::VarName,
-    vi::SimpleOrThreadSafeSimple,
-)
-    value = init(rng, dist, sampler)
-    # Transform if we're working in unconstrained space.
-    f = to_maybe_linked_internal_transform(vi, vn, dist)
-    value_raw, logjac = with_logabsdet_jacobian(f, value)
-    vi = BangBang.push!!(vi, vn, value_raw, dist)
-    vi = accumulate_assume!!(vi, value, -logjac, vn, dist)
-    return value, vi
-end
 
 function settrans!!(vi::SimpleVarInfo, trans)
     return settrans!!(vi, trans ? DynamicTransformation() : NoTransformation())

src/utils.jl

@@ -456,50 +456,6 @@ function recombine(d::MultivariateDistribution, val::AbstractVector, n::Int)
     return copy(reshape(val, length(d), n))
 end
 
-# Uniform random numbers with range 4 for robust initializations
-# Reference: https://mc-stan.org/docs/2_19/reference-manual/initialization.html
-randrealuni(rng::Random.AbstractRNG) = 4 * rand(rng) - 2
-randrealuni(rng::Random.AbstractRNG, args...) = 4 .* rand(rng, args...) .- 2
-
-istransformable(dist) = link_transform(dist) !== identity
-
-#################################
-# Single-sample initialisations #
-#################################
-
-inittrans(rng, dist::UnivariateDistribution) = Bijectors.invlink(dist, randrealuni(rng))
-function inittrans(rng, dist::MultivariateDistribution)
-    # Get the length of the unconstrained vector
-    b = link_transform(dist)
-    d = Bijectors.output_length(b, length(dist))
-    return Bijectors.invlink(dist, randrealuni(rng, d))
-end
-function inittrans(rng, dist::MatrixDistribution)
-    # Get the size of the unconstrained vector
-    b = link_transform(dist)
-    sz = Bijectors.output_size(b, size(dist))
-    return Bijectors.invlink(dist, randrealuni(rng, sz...))
-end
-function inittrans(rng, dist::Distribution{CholeskyVariate})
-    # Get the size of the unconstrained vector
-    b = link_transform(dist)
-    sz = Bijectors.output_size(b, size(dist))
-    return Bijectors.invlink(dist, randrealuni(rng, sz...))
-end
-################################
-# Multi-sample initialisations #
-################################
-
-function inittrans(rng, dist::UnivariateDistribution, n::Int)
-    return Bijectors.invlink(dist, randrealuni(rng, n))
-end
-function inittrans(rng, dist::MultivariateDistribution, n::Int)
-    return Bijectors.invlink(dist, randrealuni(rng, size(dist)[1], n))
-end
-function inittrans(rng, dist::MatrixDistribution, n::Int)
-    return Bijectors.invlink(dist, [randrealuni(rng, size(dist)...) for _ in 1:n])
-end
-
 #######################
 # Convenience methods #
 #######################