Replace old Gibbs sampler with the experimental one.

[mhauru]

Closes #2318.

Work in progress.

[codecov]

Codecov Report

Attention: Patch coverage is 68.94410% with 50 lines in your changes missing coverage. Please review.

Project coverage is 82.32%. Comparing base (78d110a) to head (6f9679a).

Files with missing lines	Patch %	Lines
src/mcmc/gibbs.jl	68.55%	50 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2328      +/-   ##
==========================================
- Coverage   86.80%   82.32%   -4.49%     
==========================================
  Files          24       21       -3     
  Lines        1599     1533      -66     
==========================================
- Hits         1388     1262     -126     
- Misses        211      271      +60     

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

[coveralls]

Pull Request Test Coverage Report for Build 11016717132

Details

• 111 of 161 (68.94%) changed or added relevant lines in 2 files are covered.
• 25 unchanged lines in 4 files lost coverage.
• Overall coverage decreased (-4.5%) to 82.538%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/mcmc/gibbs.jl	109	159	68.55%

Files with Coverage Reduction	New Missed Lines	%
src/mcmc/Inference.jl	1	87.73%
src/mcmc/ess.jl	1	96.36%
src/mcmc/particle_mcmc.jl	5	91.22%
src/mcmc/abstractmcmc.jl	18	63.79%

Totals
Change from base Build 10703372980:	-4.5%
Covered Lines:	1262
Relevant Lines:	1529

---

💛 - Coveralls

[yebai]

Gibbs(@varname(x) => hmc, @varname(x) => hmc, @varname(y) => mh)

This looks rather awkward. Can we introduce a simple wrapper, Repeated and support:

Gibbs(@varname(x) => Repeated(hmc, n), @varname(y) => mh)

[mhauru]

Yes. We had a chat about a closely related issue with @torfjelde too, I'll rework the interface around this a bit.

[mhauru]

@torfjelde, if you have a moment to take a look at the one remaining test failure, would be interested in your thoughts. We are sampling for a model with two vector variables, m and z, and we seem to somehow end up with a case where there's a VarInfo with only z in it, but the sampler is looking for m too. I wonder if it's something about the interaction between particle sampling with Libtask and how the new Gibbs does things with the local varinfos. The test that fails is this one:

    @testset "dynamic model" begin
        @model function imm(y, alpha, ::Type{M}=Vector{Float64}) where {M}
            N = length(y)
            rpm = DirichletProcess(alpha)

            z = zeros(Int, N)
            cluster_counts = zeros(Int, N)
            fill!(cluster_counts, 0)

            for i in 1:N
                z[i] ~ ChineseRestaurantProcess(rpm, cluster_counts)
                cluster_counts[z[i]] += 1
            end

            Kmax = findlast(!iszero, cluster_counts)
            m = M(undef, Kmax)
            for k in 1:Kmax
                m[k] ~ Normal(1.0, 1.0)
            end
        end
        model = imm(Random.randn(100), 1.0)
        # https://github.com/TuringLang/Turing.jl/issues/1725
        # sample(model, Gibbs(MH(:z), HMC(0.01, 4, :m)), 100);
        sample(model, Gibbs(; z=PG(10), m=HMC(0.01, 4; adtype=adbackend)), 100)
    end

[torfjelde]

Will have a look at this in a bit @mhauru (just need to do some grocery shopping 😬 )

[mhauru]

Collecting links to old relevant PRs so I don't have to look for them again: #2231, #2099

[torfjelde]

Think I found the error: if the number of m increases, say, from length(m) = 2 to length(m) = 3 during the PG step, then the lines

Turing.jl/src/mcmc/gibbs.jl

Lines 57 to 65 in 6f9679a

	if has_conditioned_gibbs(context, vn)
	value = get_conditioned_gibbs(context, vn)
	return value, logpdf(right, value), vi
	end
	# Otherwise, falls back to the default behavior.
	return DynamicPPL.tilde_assume(
	rng, DynamicPPL.childcontext(context), sampler, right, vn, vi
	)

doesn't hit the Gibbs branch since @varname(m[3]) is not present in the GibbsContext 😕

[torfjelde]

doesn't hit the Gibbs branch since @varname(m[3]) is not present in the GibbsContext

I'm a bit uncertain how we should best handle this @yebai @mhauru

The first partially viable idea that comes to mind is to subset the varinfo to make sure that it only contains the correct variables. If we do this, then m[3] will just be "ignored" (in the varinfos) until we're actually sampling the m variables, in which case it would be captured correctly.

But this would not quite be equivalent to the current implementation of Gibbs, which, AFAIK, keeps the very first occurence of m around rather than resampling everytime. And naively, I would expect this to be incorrect.

Another way is to explicitly add the varinfos to the GibbsContext itself, and then, when we encounter a value that should in fact go into a different varinfo, we add it there. But this has a few issues:

1. Requires the VarInfo to be mutable.
2. Requires the VarInfo to have a container that can keep the new incoming value m[3].
3. Implementation of Gibbs does end up being more complicated than the current approach. However, it might be worth it.

Thoughts?

[yebai]

Another way is to explicitly add the varinfos to the GibbsContext itself, and then, when we encounter a value that should in fact go into a different varinfo, we add it there. But this has a few issues:

Requires the VarInfo to be mutable.
Requires the VarInfo to have a container that can keep the new incoming value m[3].
Implementation of Gibbs does end up being more complicated than the current approach. However, it might be worth it.

I lean towards the above approach and (maybe later) provide explicit APIs to inference algorithms. This will enable us to handle reversible jumps (varying model dimensions) in MCMC more flexibly. At the moment, this is only possible in particle Gibbs; if it happens in HMC/MH, inference will likely fail (silently)

EDIT: we can keep VarInfos immutable by default, and requires inference developers to hook into specific APIs to mutate VarInfos.

[torfjelde]

This does however complicate the new Gibbs sampling procedure quite drastically 😕

And it makes me bring up a question I really didn't think I'd be asking: is it then actually preferable to the current Gibbs with keeping it all in a single VarInfo with a flag to specify whether it should be sampled or not? 😬

I guess we should first have a go at implementing this for the new Gibbs and then we can see 👍

Another point to add to the conversation that @mhauru brought to my attention the other day: we also want to support stuff like Gibbs(@varname(m) => NUTS(), @varname(m) => HMC()), i.e. multiple samplers targeting the same variables. This adds a few "complications" (beyond addressing the growing model problem discussed above):

1. Need to determine which varinfo to pick from varinfos based on the varnames present / targeted.
2. A naive implementation will result in duplicated entries in varinfos. We can however address this if we really feel like it's worth it, so probably a non-issue atm.

So all in all, immediate things we need to address with Gibbs:

1. Support changing dimensions.
2. Support picking a varinfo to condition on based on the varnames present rather than based on ===.