rogue idea: make `SamplingContext` a leaf context + make a new `InitialisationContext`

[penelopeysm]

Why leaf context?

The main thing that made me come up with this was the observation that nested SamplingContext makes no sense semantically. You're either sampling with a sampler, or you're not -- you can't be sampling with two samplers.

Extra rambly context

Specifically, it was while working on TuringLang/Turing.jl#2588: in that PR, I had to make sure that any LogDensityFunction passed to sample is equipped with a SamplingContext before sending it down the AbstractMCMC pipeline. The easiest way to do this is to just wrap the existing context in a new SamplingContext, but what happens if the existing context already has a SamplingContext? Surely this should replace the existing SamplingContext (preferable), or error (less preferable).

It is technically possible to keep SamplingContext as a parent context while forbidding nesting (e.g. with an inner constructor), but that's just an objectively worse data structure.

Once we remove PriorContext and LikelihoodContext (technically done already, just needs to be released when ready), the only other leaf context would be DynamicTransformationContext. There is no world where we use DynamicTransformationContext together with SamplingContext. In other words they are mutually exclusive:

This function ...	...calls evaluate!! with this context
(e.g.) logjoint	DefaultContext
sample!! (pending #952)	SamplingContext
link!!	DynamicTransformationContext

In other words, leaf contexts can actually be used to distinguish what we are doing with a model. (See also: #510, #953, for related discussions on this idea of using contexts to figure out what operations are being carried out.)

Why InitialisationContext?

The other part of this proposal, which I think might be more controversial, is to introduce a new leaf context, i.e. InitialisationContext. The rough idea is that, just like how SamplingContext carries a sampler with it, InitialisationContext will carry an initialisation strategy with it. That is to say, I would like something like this:

abstract type InitStrategy

# replacement for SampleFromPrior
struct Prior <: InitStrategy end

# replacement for SampleFromUniform
struct UniformLinked <: InitStrategy
    lower::Real
    upper::Real
end

struct InitContext{R<:Random.AbstractRNG,I<:InitStrategy} <: AbstractContext
    rng::R
    strategy::I
end

# I'm partly using `Init` instead of `Initialisation` because it's much shorter, but
# it also avoids the AmE/BrE spelling difference

One might ask why create a new leaf context? After all, right now we are using SamplingContext(rng, SampleFromPrior(), ...) to initialise values in a VarInfo just fine.

One nice benefit is that having a specific InitStrategy struct allows us to expose, at a high level, different ways of initialising values in a VarInfo. See TuringLang/Turing.jl#2476 (comment) -- I think, the implementation of that would readily present itself if we had something like this.

I think the main difference though is to clarify the behaviour when coming across a variable that isn't yet initialised. The point would be that SamplingContext is exclusively used to allow for samplers to override the tilde-pipeline, whereas if you want to add new values to a VarInfo, you would have to use InitialisationContext. Right now, these two types of behaviour are all smushed together in the implementation for assume(..., ::SampleFromPrior, ...):

DynamicPPL.jl/src/context_implementations.jl

Lines 212 to 260 in 0810e14

	# 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!
	# Also, if we use !! we shouldn't ignore the return value.
	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)
	push!!(vi, vn, f(r), dist)
	# By default `push!!` sets the transformed flag to `false`.
	settrans!!(vi, true, vn)
	else
	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)
	return r, logpdf(dist, r) - logjac, vi
	end
	# default fallback (used e.g. by `SampleFromPrior` and `SampleUniform`)
	observe(sampler::AbstractSampler, right, left, vi) = observe(right, left, vi)
	function observe(right::Distribution, left, vi)
	increment_num_produce!(vi)
	return Distributions.loglikelihood(right, left), vi
	end

And it doesn't help that quite a number of actual samplers simply delegate to SampleFromPrior (grep for this in Turing) and so the behaviour is quite difficult to predict.

Having these two separate would allow for a cleaner separation of concerns.

Potential issue 1: when the concerns aren't separate

The immediate concern I had with this is that sometimes there are uninitialised variables when sampling. For example:

@model function patho()
    x ~ Normal()
    if x > 0
        y ~ Normal()
    else
        z ~ Normal()
    end
end

It would be very possible to create a VarInfo that has x and y, and then on a second evaluation come across z, in which case we would have to decide how to deal with it.

The truth though is that this model currently errors:

julia> sample(patho(), NUTS(), 100)
ERROR: type NamedTuple has no field z [...]

julia> sample(patho(), MH(), 100)
ERROR: type NamedTuple has no field y [...]

so we aren't immediately losing any functionality, and indeed I think that having separation of concerns might help us implement this (cursed) functionality in a better way, if it's something we actually want to do.

The other potential issue

To the best of my knowledge, SamplingContext doesn't have any nasty interactions with other contexts, in the sense that (e.g.) I don't think it ever makes a difference whether SamplingContext is nested inside ConditionContext or the other way around. Gibbs, as always, is the one area where I cannot say this with confidence.

[penelopeysm]

@mhauru thoughts welcome (and I know you just got rid of leaf contexts 😭, but if we make new leaf contexts here it comes with the benefit of simplifying parent contexts, so it's not a strict increase in complexity.)

[mhauru]

A couple of immediate thoughts:

• Would there be a difference between SamplingContext with prior sampler and InitContextwith Prior strategy?
• Regardless of whether this is the solution or not, the weird logic around SampleFromPrior in context_implementations that you linked above needs to go. It's one of the most confusing bits of code in DPPL.
• Regarding the patho example, this model is fine to sample with MH and PG (but not NUTS):

julia> @model function f(::Type{T}=Float64) where {T}
           x ~ Normal()
           if x > 0
               y = Vector{T}(undef, 2)
           else
               y = Vector{T}(undef, 3)
           end
           for i in eachindex(y)
               y[i] ~ Normal()
           end
           return nothing
       end

I'm not sure how this would interact with InitContext.

[penelopeysm]

Would there be a difference between SamplingContext with prior sampler and InitContext with Prior strategy?

I think the former would internally use the latter. Actually, that would be the same as what we are doing now: Turing's Prior calls evaluate!!(model, VarInfo(), SamplingContext(rng, SampleFromPrior))). This only works because we pass in an empty varinfo (which is what allows SampleFromPrior to generate fresh values, in that code block above). But it would be clearer.

I'll have to think a bit more about that model 🫠.

[penelopeysm]

So I can't say I fully like it, but I think the cleanest way we can handle this is still to have InitContext, and be explicit about it by doing something like this:

# This function returns an InitStrategy that the sampler uses to generate
# new values for yet-unseen variables. We could default to Prior
# Note that this already exists, it's called DynamicPPL.initialsampler.
function init_strategy(::AbstractSampler) end

function tilde_assume(ctx::SamplingContext, right, vn, vi)
    if haskey(vn, vi)
        # Individual samplers will still have to overload this method. Notice
        # that we don't need to dispatch on ctx.context because SamplingContext
        # will be a leaf context. This means that we don't have to handle all
        # possible dispatches for tilde_assume(rng, ::AbstractContext, ...).
        tilde_assume(ctx.rng, ctx.sampler, right, vn, vi)
    else
        tilde_assume(InitContext(ctx.rng, init_strategy(ctx.sampler), right, vn, vi)
    end
end

function tilde_assume(ctx::InitContext, right, vn, vi)
    # push to vi here, or if the variable exists, overwrite it
end

# default option for this method is to do nothing special, but samplers
# can override it if they need to
function tilde_assume(rng, spl::AbstractSampler, right, vn, vi)
    tilde_assume(DefaultContext(), right, vn, vi)
end

I guess a lot of this is already implicit in the SampleFromPrior method and I reckon this would be a much cleaner way of expressing the same logic. So when I say I don't fully like it, I think it's more of me expressing the sentiment that I don't like that we need to do this, rather than me disliking the proposed solution.

The main thing that makes me unhappy is that it would mean that one can just use SamplingContext to initialise new values, even though maybe one shouldn't -- for example, if init_strategy(::MH) = Prior(), then it turns out that the following

_, vi = evaluate!!(model, VarInfo(), InitContext(rng, Prior())

_, vi = evaluate!!(model, VarInfo(), SamplingContext(rng, MH())

would be equivalent, although the first one is obviously much clearer to read. I think, by separating the two types of context, I really wanted to prevent the use of the latter. But I suppose having two separate things still helps, because we can then say for sure that

• InitContext will always add new values
• SamplingContext only adds new values if it needs to

and so if the code is written properly (e.g. picking the InitContext over the SamplingContext above) it's easier to intuit the meaning of code when reading it. Plus, docstrings, of course -- we need docstrings and stuff. The whole tilde-pipeline is as good as undocumented.

[mhauru]

I think I can get onboard with SamplingContext being a leaf.

I'm less certain about InitContext. I find it semantically a bit confusing. The very meaning of sampling is to get new values for the model's variables. Most of serious MCMC samplers do a thing where they use some process to generate candidate values and then only use model evaluations to check whether those values are good, which is why evaluating with a SamplingContext for most samplers doesn't change the values in vi. But some samplers, like SampleFromPrior, do, and semantically that makes abundant sense. I'm imagining being new to the codebase. First I meet SamplingContext and go "got it, this is what you use when you're using a sampler to sample a new set of values". Then I meet InitContext and am confused about why there's a distinct context that is explicitly about putting newly-sampled values into vi. (The fact that SamplingContext in most instances seemed to just be evaluating logp with some extra bells or whistles confused me a lot when I was new to the codebase.)

I think I like the init_strategy(::AbstractSampler) construction, because having to generate values for new variables mid-sampling is a problem that affects all samplers (the need for it is defined by the model). However, it seems to me that this could be implemented as naturally with something like

function tilde_assume!!(ctx::SamplingContext, right, vn, vi)
    if !haskey(vn, vi)
        vi = tilde_assume!!(SamplingContext(init_strategy(ctx.sampler)), right, vn, vi)
    end
    # business as usual
end

(I'm also unsure whether this logic should be implemented in Turing or if it should be the responsibility of each sampler to decide what to do, but that's a distinct question.)

I'm not adamantly convinced that I don't want InitContext, but so far I am not convinced that I do.

[penelopeysm]

Most of serious MCMC samplers do a thing where they use some process to generate candidate values and then only use model evaluations to check whether those values are good, which is why evaluating with a SamplingContext for most samplers doesn't change the values in vi. But some samplers, like SampleFromPrior, do, and semantically that makes abundant sense.

"Most samplers don't need to change the values, but SampleFromPrior does" is precisely why I think SampleFromPrior should not be a sampler. There's a very subtle distinction between

• evaluating the feasibility of a new set of values generated by some other process (e.g. leapfrog integration in HMC), vs
• sampling a new set of values through model evaluation

SamplingContext{<:Hamiltonian} does the former, not the latter. (In this regard, I'd argue the name is a misnomer, but I wouldn't go as far as renaming it.)
SamplingContext{<:SampleFromPrior} does the latter*, and my point is really that I don't like that they are lumped under the same umbrella

(I recognise you said all this, but surely that proves my point? Why do we have two different kinds of samplers, surely it would be better to express this using the type system rather than docstrings?)

* Well, actually SampleFromPrior doesn't sample new values. In this regard it actually behaves more like DefaultContext. That's why many samplers fall back on it, and that's why in src/mcmc/ess.jl there's this ridiculous need to set del flags in the VarInfo. I legitimately think this is a mistake (historical baggage) and samplers should directly fall back on DefaultContext. SampleFromPrior should be changed to always set fresh values even if they already exist in the varinfo.

alternative

The main compromise I guess is to keep SampleFromPrior and SampleFromUniform as actual samplers and keep using SamplingContext for them. But because we want to be able to let people specify initialisation strategies, not samplers (it doesn't make any sense to "initialise a varinfo with HMC"!) we can still use a different subtype:

abstract type InitStrategy <: AbstractSampler end

function init_strategy(::AbstractSampler) end

struct SFPrior <: InitStrategy end
struct SFUniformLinked{T<:Real} <: InitStrategy
    lower::T
    upper::T
end

struct SamplingContext{R<:AbstractRNG,S<:AbstractSampler} <: AbstractContext
    rng::R
    sampler::S
end

The main difference is that what I'm proposing as InitStrategy will still be an AbstractSampler (by subtyping), and the two contexts will be merged into one. (Or kept as one.) Then tilde_assume(::SamplingContext) will be as you wrote.

I'd still much prefer having them separate, though.

[mhauru]

There's a very subtle distinction between

• evaluating the feasibility of a new set of values generated by some other process (e.g. leapfrog integration in HMC), vs
• sampling a new set of values through model evaluation

SamplingContext{<:Hamiltonian} does the former, not the latter. (In this regard, I'd argue the name is a misnomer, but I wouldn't go as far as renaming it.)
SamplingContext{<:SampleFromPrior} does the latter*, and my point is really that I don't like that they are lumped under the same umbrella

(I recognise you said all this, but surely that proves my point? Why do we have two different kinds of samplers, surely it would be better to express this using the type system rather than docstrings?)

But what is the meaning of SamplingContext in your proposal? "The context that is used when a sampler is evaluating the model but without sampling new values." That's not a very graspable concept, it makes the logic of the codebase harder to follow. Also, names matter a lot for code understandability.

I'm not happy with the confusing merging of the two sorts of evaluate calls that samplers need to do either. For instance, I agree that

SampleFromPrior should be changed to always set fresh values even if they already exist in the varinfo.

and we should get rid of "del" in VarInfo. But I'm struggling to come up with the right concepts and names to make this part of the codebase intuitive. I'm not sure SamplingContext vs InitContext is the solution.

I think what I would really want is to have all the cases where a sampler is just evaluating the model as part of its sampling process, but not actually doing any sampling during the evaluation, to be called something other than SamplingContext. Then SamplingContext could be doing what you call InitContext, which would naming-wise make more sense to me. However, I don't know how we could handle all the other stuff that samplers need to do in the tilde-pipeline. Maybe I should take a look and remind myself what all that "other stuff" is.

So I went and had a look, and there actually isn't that much "other stuff". With a quick skim, the only cases in Turing.jl where a SamplingContext{<:MyCoolSampler} was actually doing something other than sampling from the prior I think were Gibbs and particle methods. I think they could be justified in implementing their own context since they do proper weird stuff; Gibbs already does. Could we then say that SamplingContext will always have to do actual sampling of each variable it meets, replacing the old value (like prior and uniform), and if a sampler really needs to evaluate the model and make it do something other than the usual stuff that DPPL provides then that sampler needs to define its own context?

[penelopeysm]

We discussed on Slack, so just repeating here:

Could we then say that SamplingContext will always have to do actual sampling of each variable it meets, replacing the old value (like prior and uniform), and if a sampler really needs to evaluate the model and make it do something other than the usual stuff that DPPL provides then that sampler needs to define its own context?

I wouldn't be opposed to that, as long as we disambiguate the two current cases of SamplingContext.

I agree many SamplingContexts aren't really needed. (I looked at ESS, for example, and I don't think we need a SamplingContext{ESS}. It exists for two reasons -- one is to get fresh variables for the prior, which we don't need a new context for, and one is to calculate the likelihood, which we can soon do with getloglikelihood. So there's no real reason why that needs to exist.)

As I understand it, our proposals aren't too different, and I'm not super duper fussed about which one we go for.

Current	What this does	Penny thinks this should be	Markus thinks this should be
SamplingContext(SampleFromPrior)	If variables are absent, gets new ones from prior	InitContext(Prior)	SamplingContext(Prior)
SamplingContext(HMC)	Evaluates log-density during HMC	SamplingContext(HMC)	DefaultContext
SamplingContext(PG)	Evaluates log-density during PG	SamplingContext(PG)	PGContext

We agreed that I'll raise this at one of the Monday meetings, but @yebai maybe you also had a view on this.

[yebai]

I find the name SamplingContext uninformative and would suggest taking this opportunity to eliminate it. My preferred approach is a reconciliation of the two excellent proposals by Penny and Markus.

Current	What this does	Hong thinks this should be
SamplingContext(SampleFromPrior)	If variables are absent, gets new ones from prior	InitContext
SamplingContext(HMC)	Evaluates log-density during HMC	DefaultContext
SamplingContext(PG)	Evaluates log-density during PG	PGContext

[penelopeysm]

I didn't think about this earlier, but it strikes me now that we still need a way to indicate that 'we are in the middle of sampling and thus it's alright to add a new variable if it doesn't exist'. With DefaultContext, trying to evaluate with a varinfo that isn't complete will error (which is probably the right behaviour), whereas with our current SamplingContext we say that's alright (because we want to handle dynamic models). So there is still a difference 😢

Of course there are many ways we can express this that don't involve keeping SamplingContext -- for example, we could have a type parameter on DefaultContext -- but I guess the question is what's the best way.

TLDR:

• DefaultContext now: can never add new variables. We are quite sure we want to keep this.
• SamplingContext now: (usually) doesn't overwrite existing variables, but can add new variables when needed. Not sure what to do with this.
• InitContext in future: will always add / overwrite existing variables. We seem quite sure we want to add this.

Anyway, I'll put this on the agenda for Monday's meeting and we can discuss in more detail then. I'll also put together an initial PR (not sure when, maybe next week, or the week after) to see what are the effects of doing this + downstream effects on Turing.