Remove Sampler, remove InferenceAlgorithm, transfer initialstep, init_strategy, and other functions from DynamicPPL to Turing

[penelopeysm]

TuringLang/DynamicPPL.jl#1037 deletes DynamicPPL.Sampler and its interface. This PR makes the necessary upstream adjustments.

Once this PR and the DPPL counterpart are merged, all of the MCMC sampling behaviour will be self-contained in Turing: this means that we can easily refactor sampler code in Turing without being held back by interface code in DynamicPPL. This will finally pave the way for things like sampling with LogDensityFunction #2555 and allow long-standing things to be fixed like weird parameters to NUTS #2678. It's been a few months coming so I am very pleased with this!

Literally all of the actual code changes in this PR are perfunctory boilerplate, like copy-pastes from DynamicPPL, changing x::Sampler{<:X} to x::X, etc. However, there is one conceptual bit that I think is worth explaining:

Special behaviour when sampling from Turing models with Turing samplers

As discussed in, e.g. #2413, there are certain special behaviours that are enabled when somebody calls sample(model, NUTS(), 1000). For example:

• check_model runs checks on the model,
• initial_params is set to DynamicPPL.init_strategy(spl) (now Turing.Inference.init_strategy(spl)),
• chain_type is set to MCMCChains.Chains, ...

Previously the way this was accomplished was by

1. declaring NUTS <: InferenceAlgorithm,
2. sample(rng, ::AbstractModel, ::InferenceAlgorithm, N) would then wrap NUTS in DynamicPPL.Sampler
3. sample(rng, ::AbstractModel, ::DynamicPPL.Sampler, N) would then do all the special behaviour

This PR changes it such that the special behaviour is triggered in sample(rng, ::DynamicPPL.Model, ::AbstractSampler), in src/mcmc/abstractmcmc.jl. That is to say, NUTS is no longer 'a special class of sampler'. If you write your own sampler that knows how to handle Turing models (i.e. if you define step(rng, ::DynamicPPL.Model, ::MySampler)), then you will get all the special behaviour when you call sample(model, MyNewSampler(), 1000).

This approach is a lot simpler than what came before, but as pointed out in #2413 there are some concerns with method ambiguities. For example, if somebody defines

sample(::AbstractModel, ::MyNewSampler, N) = ...

then calling sample(dynamicppl_model, MyNewSampler(), 1000) will lead to method ambiguities. However, it's been almost a year since then, and there are a number of reasons why I'm not really fussed about the ambiguities and why I no longer agree with the arguments in that thread (indeed I don't agree with my own comments from that thread).

I've gated the rest behind a details tag to avoid overwhelming. If you don't see a problem with the above, then feel free to ignore.

Method ambiguities

1. As explained in Remove (duplicate) samplers being defined explicitly in Turing.jl #2413 the only meaningful case where somebody might write such a sample method is if their sampler was a 'meta-sampler', i.e., it works with all models without knowing anything about the internal structure of said models. In practice this cannot be done because AbstractModel has no interface and if you don't know anything about the model, you can't do anything meaningful with it. Look at Gibbs, for example: in principle it shouldn't need to know anything about the model because it just shuttles it between component samplers. In practice you have to specialise on the model type and include a ton of stuff like setparams_varinfo just to make it even work.

2. Method ambiguities have a very straightforward solution which is to declaring an extra, more specific method. I think that is a perfectly fine price to pay if you are writing a (genuine) meta-sampler and want it to work with Turing models specifically. For example this is the case with RepeatSampler, it's just five extra lines of code to make it work.

3. Note that in the original proposal in Remove (duplicate) samplers being defined explicitly in Turing.jl #2413 the method ambiguities would be avoided but if somebody wrote a model and wanted to hook into the special behaviour, they would have to write the more specific method anyway! So IMO this is no worse.

4. I haven't exactly seen an explosion of activity in the Julia ecosystem around creating meta-samplers, so I'm quite unconvinced that in practice this will cause any problems

5. I think the way the code is now written (i.e. in this PR) is far closer to what is intended. Special behaviour is tied to the fact that you are sampling a DynamicPPL.Model which brings with it a rich structure that lets you check its validity, lets you group its outputs into chains, etc. Previously, special behaviour was tied to the fact that "it was a sampler defined in Turing.jl", which was not a very meaningful metric for association.

[github-actions]

Turing.jl documentation for PR #2689 is available at:
https://TuringLang.github.io/Turing.jl/previews/PR2689/

[codecov]

Codecov Report

❌ Patch coverage is 81.37255% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.34%. Comparing base (deff3fd) to head (593af43).

Files with missing lines	Patch %	Lines
src/mcmc/abstractmcmc.jl	53.33%	14 Missing ⚠️
src/mcmc/ess.jl	50.00%	2 Missing ⚠️
src/mcmc/gibbs.jl	84.61%	2 Missing ⚠️
src/mcmc/Inference.jl	0.00%	1 Missing ⚠️

Additional details and impacted files

@@                Coverage Diff                @@
##           py/dppl-0.38    #2689       +/-   ##
=================================================
+ Coverage         57.89%   84.34%   +26.45%     
=================================================
  Files                22       21        -1     
  Lines              1387     1412       +25     
=================================================
+ Hits                803     1191      +388     
+ Misses              584      221      -363     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[penelopeysm]

For full transparency, I will confess that I had to add two methods to remove method ambiguities caused by this PR. This is the first one.

Both arise from the default definition of AbstractMCMC.step

Turing.jl/src/mcmc/abstractmcmc.jl

Lines 141 to 163 in b5e0760

	# TODO(penelopeysm): Remove initialstep and generalise MCMC sampling procedures
	function initialstep end
	function AbstractMCMC.step(
	rng::Random.AbstractRNG,
	model::DynamicPPL.Model,
	spl::AbstractSampler;
	initial_params,
	kwargs...,
	)
	# Generate the default varinfo. Note that any parameters inside this varinfo
	# will be immediately overwritten by the next call to `init!!`.
	vi = default_varinfo(rng, model, spl)
	# Fill it with initial parameters. Note that, if `InitFromParams` is used, the
	# parameters provided must be in unlinked space (when inserted into the
	# varinfo, they will be adjusted to match the linking status of the
	# varinfo).
	_, vi = DynamicPPL.init!!(rng, model, vi, initial_params)
	# Call the actual function that does the first step.
	return initialstep(rng, model, spl, vi; initial_params, kwargs...)
	end

which I would like to remove (because it doesn't allow for warmup), but I don't want to do it in this PR / version.

[penelopeysm]

This is the other one.