`forkr` is likely broken for particle Gibbs

[yebai]

forkr is likely broken without explicit dependence on the replaying mechanism in Turing.jl. This is required to correctly implement the reference particle for particle Gibbs. A potential fix is the switch to a simpler replying mechanism by fixing the random seed for the reference particle.

[devmotion]

A potential fix is the switch to a simpler replying mechanism by fixing the random seed for the reference particle.

I thought a bit more about this and noticed that the main problem will be the resampling step - even if the random seed and the random number generator that is passed to the model execution is fixed, then in the resampling step we would like to obtain diverging traces. I.e., it is impossible to fix the model execution in advance by setting a seed initially. Maybe a workaround could be to save the seed in every step and to explicitly reseed the random number generators (which would have to be separate for each particle) in every resampling step.

[yebai]

It sounds like a sensible solution. Maybe consider introducing a simple assume function, so that we have more control of the model execution than directly calling rand?

[yebai]

The following package from the Gen team might be interesting.

https://github.com/probcomp/TracedRandom.jl

Maybe we can consider introducing a customised rand / simulate function, which accepts a unique RV name and a distribution as arguments. The unique RV name is a stronger assumption than Turing but is probably fine for a pure particle MCMC package.

Given that all RVs are explicitly uniquely named, they can be stored in a standard Julia dictionary (or something similar but more performant). It would allow particle Gibbs (and its ancestor resampling variant) to be supported, without resorting the general-purpose hefty VarInfo mechanism from Turing.

[devmotion]

I don't think it's a good idea to add yet another signature to rand that is not compatible with any other use of rand I've seen. In general, it seems the package is mostly type piracy?

Somehow to me it feels the problem of replaying a trace and saving random variables in some structure are two distinct problems (and hence I am not sure if this should be two separate issues). If the samples can be obtained in the order they appear in the model, then in the most basic setting users could map them to variable names manually?

Alternatively, maybe one could perform some IR-rewriting of the model that allows to inject an "context" that saves sampled variables in some data structure.

[yebai]

I agree that a clean, minimal implementation is preferred here. We'd like to keep AdvancedPS's dependency footprint small too.

Somehow to me it feels the problem of replaying a trace and saving random variables in some structure are two distinct problems

I think we can avoid saving random variables as long as they can be replayed/reconstructed, e.g. by fixing random seeds.

As for the replaying mechanism, I am thinking about something in the lines of a TracedRNG <: AbstractRNG. This TracedRNG could record the seed, and the number of randomnesses an RNG has generated. In addition, we can have a VectorOfTracedRNG <: AbstractRNG so that it can support the type of operations required by ancestor resampling in particle Gibbs.

# Illustration code
mutable struct TracedRNG{T}  <: AbstractRNG where T<:AbstractRNG
   count::Int # Number of randomness it has generated
   seed   # Seed of RNG
   rng::T # Current state of RNG
end

VectorOfTracedRNG <: AbstractRNG is a vector of TracedRNGs. It is helpful in cases where we need to mix and match several streams of randomness, such as when forking the reference particle in particle Gibbs.

Any thoughts?

[FredericWantiez]

To be able to replay the trace given the TracedRNG, I think we need to change the Trace definition to explicitly reference the rng ? That also means we need to change the model definition to take the rng explicitly to be able to keep track of the randomness generated if we mutate the rng in rand(rng::TracedRNG, ...)

[devmotion]

Yes, I assume the model has to include the rng as well, as noted e.g. in #20 (comment) and #20 (comment). Other algorithms such as MH or HMC were updated when the algorithms in Turing were switched to AbstractMCMC 2, but this was not done for the particle methods.

Maybe it would be useful to also not allow general RNGs but use user-provided RNGs (or the implicitly provided global RNG) only to seed a special class of internally used RNGs (IIRC e.g. StochasticDiffEq does this). Possibly this could allow additional optimizations - e.g., maybe it could be useful to use a RNG design similar to Jax that is based on a counter PRNG and a splitting model? Intuitively, the splitting model could be useful when resampling particles.

[yebai]

Closed by #23