`initial_state` is not used when `resume_from` is not specified

[ysfoo]

sample does not use initial_state if specified without resume_from. EDIT: This applies when sampler is of type Sampler{<:AdaptiveHamiltonian}.

I encountered this when I tried to resume one chain after running multiple chains, and thought I'd do this using initial_state instead of resume_from, especially in light of #2675.

MWE:

using Turing, Random

@model function gdemo(x)
    s² ~ InverseGamma(2, 3)
    m ~ Normal(0, sqrt(s²))
    for i in eachindex(x)
        x[i] ~ Normal(m, sqrt(s²))
    end
end

rng = Random.default_rng();
Random.seed!(rng, 1);
data = 4 .+ 0.1 .* randn(rng, 100);

chain1 = sample(rng, gdemo(data), NUTS(), 10, nadapts=1000, save_state=true)
display(Array(chain1))

chain2 = sample(rng, gdemo(data), NUTS(), 10, nadapts=0, initial_state=chain1.info.samplerstate)
display(Array(chain2)) # initial params is drawn from prior instead of using initial_state

[penelopeysm]

display(Array(chain2))

The parameters of chain1.info.samplerstate are not in general going to be the same as the parameters of the first sample of chain2. That's because chain2 actually continues from chain1.info.samplerstate rather than returning that as the first state.

For example notice that the samples begin from 11 here rather than from 10:

julia> using AbstractMCMC

julia> struct M <: AbstractMCMC.AbstractModel end

julia> struct S <: AbstractMCMC.AbstractSampler end

julia> function AbstractMCMC.step(rng, ::M, ::S, state=0)
           return state + 1, state + 1
       end

julia> sample(M(), S(), 10; initial_state=10)
Sampling 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| Time: 0:00:00
10-element Vector{Int64}:
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20

[ysfoo]

Yep, understood. My problem is that this AbstractMCMC.sample method ignores initial_state when resume_from is nothing, so the second chain does not resume from the given initial_state.

In the example above, the traces for s² and m for chain1 are

10×2 Matrix{Float64}:
 0.186341  3.97191
 0.171828  3.94066
 0.182928  3.9681
 0.20692   3.86289
 0.24893   4.06012
 0.218585  4.03193
 0.278705  3.93844
 0.211338  4.00237
 0.226351  3.97471
 0.223113  3.99839

This is the typical set of the posterior (m is closer to 4).

But when resuming from the last state of this chain, the traces for s² and m for chain1 are

10×2 Matrix{Float64}:
 3.82191   1.33987
 4.06332   1.41501
 4.27018   1.43835
 5.95238   2.0899
 5.56198   2.09329
 5.47808   2.10791
 4.74372   2.17253
 4.85457   2.25468
 3.54586   2.98396
 0.684549  3.93981

You can see that m is drifting from the typical set of the prior towards the typical set of the posterior, which shouldn't happen since we're resuming from the typical set of the posterior.

Ideally, I think this call to AbstractMCMC.mcmcsample should pass initial_state, otherwise it seems counterintuitive to me that initial_state does nothing without specifying resume_from.

[penelopeysm]

Ahh, thank you! I didn't realise it was a NUTS-only thing.

Yup, that definitely needs initial_state to be passed through.