Better support for AbstractSampler

[torfjelde]

Now that we have DynamicPPL.LogDensityFunction + AbstractMCMC.LogDensityModel, it would be nice to drop the usage of Turing.Inference.InferenceAlgorithm and allow "arbitrary" AbstractMCMC.AbstractSampler to be used with in Turing without much effort on the user side.

This would of course also help alleviate the discrepancy in configuration that is present on the package containing the implementation of the sampler and the corresponding "inference algorithm" in Turing.jl.

The issue with just replacing InferenceAlgorithm with AbstractSampler and AbstractMCMC.AbstractModel with DynamicPPL.Model, as is done in #2008, is that of method-ambiguities. I also assume that method ambiguities were part of the reason why InferenceAlgorithm was kept around.

There are, as I see it, two solutions:

1. Replace the currently present overloads of AbstractMCMC.sample in Turing.jl by Turing's own sample, i.e. make Turing.sample !== AbstractMCMC.sample. Turing.sample then calls AbstractMCMC.sample under the hood but with some additional niceties.
2. Keep something equivalent to InferenceAlgorithm or a wrapper type and continue overloading AbstractMCMC.sample.

Personally, I'm very much in favour of option (1). This has several benefits:

1. Removes the type-piracy.
2. Maintains the status-quo.

[torfjelde]

@devmotion @yebai @JaimeRZP thoughts?

[JaimeRZP]

Hi Tor!

Thanks for opening an issue.
After some thought I am also in favour of 1; i.e. Turing.sample .
I think besides the points you have raised I can also see how it would be easier to debug in case the user manages to dispatch to the wrong function.
Also for some samplers Turing dispatches to sampler specific instance AbstractMCMC.sample. Having all of these being specified within Turing seems the most elegant solution.

What do others think?

Best,
Jaime

[JaimeRZP]

I have been playing around with these ideas and I have noticed that even if we get rid of InferenceAlgorithm would still need to wrap AbstractSampler in DynamicPPL.Sampler() if we do not wish to change many things in one single PR. Would somethign like the code below also lead to method ambiguities?

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::Model,
    sampler::Sampler{<:AbstractSampler},
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
.
.
.
end

[JaimeRZP]

Quick note: If we do this we will have the same type ambiguity that we would have in sample in bundle_samples.

[devmotion]

One consequence of (1) is that many users will run into warnings about conflicting exports and many scripts and tutorials will can't be used with newer Turing versions - AbstractMCMC.sample itself is not owned by AbstractMCMC but just StatsBase.sample. (1) might still be preferable but that's a somewhat unfortunate consequence of it, I think.

I wonder if there is a better way to fix these method ambiguities. What's the main problem to fix here? I think the problem is mainly that #2008 does keep sample definitions with generic sampler and DynamicPPL.Model - even though I think they should be removed. If a sampler package defines sample, Turing should be able to reuse it automatically without any custom definitions if this implementation is defined for generic LogDensityModels.

[torfjelde]

What's the main problem to fix here?

Problem is: how do we avoid method ambiguities while at the same time allowing Turing.sample (whatever this method is) to add Turing's bells and whistles, e.g.

• Good support for conversion into a Chains, i.e. if chain_type=Chains (which it should also be by default in Turing.sample) then we should extract parameter names, etc.
• Convert DynamicPPL.Model into a LogDensityFunction or whatever the actual sampler wants.

I think the problem is mainly that #2008 does keep sample definitions with generic sampler and DynamicPPL.Model - even though I think they should be removed.

So before that was InferenceAlgorithm + AbstractModel. I def agree that we should remove these overloads or something, but we want to do that while we preserve the current behavior of Turing.sample, right?

If a sampler package defines sample, Turing should be able to reuse it automatically without any custom definitions if this implementation is defined for generic LogDensityModels.

Which def are you referring to?

[torfjelde]

One consequence of (1) is that many users will run into warnings about conflicting exports and many scripts and tutorials will can't be used with newer Turing versions - AbstractMCMC.sample itself is not owned by AbstractMCMC but just StatsBase.sample.

Very much agree with this though 😕

[devmotion]

while at the same time allowing Turing.sample (whatever this method is) to add Turing's bells and whistles, e.g.

My premise was that these would have to be removed or integrated in AbstractMCMC (possibly not the exact default values but some way to specify them).

Convert DynamicPPL.Model into a LogDensityFunction or whatever the actual sampler wants.

If the sampler expects a LogDensityModel, shouldn't this work automatically?

[torfjelde]

My premise was that these would have to be removed or integrated in AbstractMCMC (possibly not the exact default values but some way to specify them).

But how do we do either of these while still preserving the desired behavior mentioned above? 😕

Complete removal of Turing.sample implementations seems ideal, but I don't understand how we can do that and preserve the status quo from a Turing-user's persective 😕

If the sampler expects a LogDensityModel, shouldn't this work automatically?

I don't think I fully understand 😕 Can you write the signature you want Turing.sample to have (and whether it's AbstractMCMC or not)?

[devmotion]

I assumed that Model <: AbstractMCMC.LogDensityModel. Isn't that what we want? Then sample(rng, ::LogDensityModel, ::MyCustomSampler) would be sufficient (modulo Turing-specific defaults). But it seems currently it's just a subtype of AbstractMCMC.AbstractModel (https://github.com/TuringLang/AbstractPPL.jl/blob/0f289206b22da5feee03218c157afb28706ed8ec/src/abstractprobprog.jl#L11).

[torfjelde]

I assumed that Model <: AbstractMCMC.LogDensityModel. Isn't that what we want?

Uhmm that is not at all what I hand in mind 😅 But I don't quite see how this helps. Sure, it does let us remove the sample functions from Turing, but we're also then losing everything that comes with that.

Then sample(rng, ::LogDensityModel, ::MyCustomSampler) would be sufficient (modulo Turing-specific defaults).

But the Turing-speciific stuff is sort of the important bit, no? 😅 Otherwise, yes things become much simpler.

[devmotion]

My feeling was that they can easily be kept by changing the upstream code a bit. After all, it seems the only Turing-specific keyword arguments are chain_type = Chains and resume_from? Support for the latter could be added in AbstractMCMC I think, and the chain_type default could be made a function of the model (and sampler, I assume? - but then we need clear instructions for who is allowed to implement it to avoid method ambiguities, I guess).

[torfjelde]

But how do we extract the variable names in this scenario? This also needs to enter somewhere in the process, right?

chain_type default could be made a function of the model (and sampler, I assume? - but then we need clear instructions for who is allowed to implement it to avoid method ambiguities, I guess).

But would that help us in this case? We would dispatch on ::Model and ::AbstractSampler then, which would still leave room for ambiguity?

[devmotion]

The variable names? Isn't that done by save!! and bundle_samples?

We can also prescribe the priority of arguments by defining eg

default_chain_type(model, sampler) = default_chain_type(sampler)
default_chain_type(sampler) = Any

and demanding that sampler packages only implement default_chain_type(::MySampler).

But, honestly, maybe it would be best to remove chain_type completely. I've never liked it very much, it doesn't seem a good approach for achieving composability. Maybe instead we should just demand that sample returns something that satisfies a Chains-interface (similar to Tables.jl, with some defaults for base types like Vector{<:NamedTuple}). And then Chains pacakges could define their constructors/converters for anything that satisfies this interface.

[torfjelde]

The variable names? Isn't that done by save!! and bundle_samples?

Aye, but on what do we dispatch bundle_samples if we remove InferenceAlgorithm?

But, honestly, maybe it would be best to remove chain_type completely. I've never liked it very much, it doesn't seem a good approach for achieving composability. Maybe instead we should just demand that sample returns something that satisfies a Chains-interface (similar to Tables.jl, with some defaults for base types like Vector{<:NamedTuple}). And then Chains pacakges could define their constructors/converters for anything that satisfies this interface.

Very happy to go down this route! Though it will require some more effort on our part. I do agree with the comment on chains_type too.

EDIT: The effort-comment isn't to discourage the idea, just to comment on the effort vs. the above proposals. But as we've said, the above proposals are much worse alterantives.