Decide the fate of `VarInfo.num_produce`

[mhauru]

This is a feature of VarInfo that's only used by particle samplers. Having it stored in VarInfo doesn't feel like the right place, and we would like to get rid of it when we refactor VarInfo/SimpleVarInfo. However, having the sampler keep track of it is apparently not easy either, and might require changes to how some contexts work around samplers. We could also drop it, but that would mean losing existing particle sampler functionality.

Discussed with @torfjelde, @sunxd3, and @willtebbutt.

[torfjelde]

@yebai should also probably have a look at this :+1:

[yebai]

We could move this to AdvancedPS, e.g. adding a num_produce field to AdvancedPS.LibtaskModel:

https://github.com/TuringLang/AdvancedPS.jl/blob/6c2cee380ce53ab25483e57aaaf381fbb8f31fd3/src/model.jl#L49-L53

Each LibtaskModel realisation corresponds to a particle / VarInfo.

cc @FredericWantiez @willtebbutt

[yebai]

@mhauru, do you want to try this? I'm happy to help you through the process if you have any questions.

[yebai]

For reference, the only place that we use the num_produce field is below:

https://github.com/TuringLang/DynamicPPL.jl/blob/24a73809b1a9f0f5b4e5f907f737b57c6eaf801a/src/varinfo.jl#L1716-L1719

which is called by a function in Turing,

https://github.com/TuringLang/Turing.jl/blob/78d110afc17daaedbc492695d258c5eb8541463e/src/essential/container.jl#L41-L44

Particle MCMC samplers use this to remove obsolete random variables in VarInfo after a particle is copied/forked (as part of the resampling step in particle filtering algorithms).

[mhauru]

Can do. Will take me a moment though, I'll have to start by reading some wikipedia to get a basic idea of what particle MCMC does.

[yebai]

I thought I'd write more about num_produce if it is helpful.

The num_produce variable is used by particle Gibbs (PG), specifically during the "particle copying" step after each resample step. In PG, each resample step is triggered by a new observation (in state-space models, we filter upon new observations). For "particle copying", we need to replay the randomness stored in VarInfo. But we should terminate this replaying process at specific points of model execution (if you know the PG sampler, this is due to PG's reference particle). As such, the num_produce variable is used to determine these specific points to stop replaying.

In the DynamicPPL code, we delete random variable i with orders[i] greater than num_produce (see the above-referenced code block). Later, the tilde pipeline will resample random variables if they cannot be found in VarInfo, thus equivalent to stopping relaying.

[mhauru]

There seem to be a couple of other spots where num_produce is accessed:

• When push!!ing a new variable into a VarInfo: https://github.com/TuringLang/DynamicPPL.jl/blob/24a73809b1a9f0f5b4e5f907f737b57c6eaf801a/src/varinfo.jl#L1627
• In assume when using SampleFromUniform, and similarly in get_and_set_val!: https://github.com/TuringLang/DynamicPPL.jl/blob/24a73809b1a9f0f5b4e5f907f737b57c6eaf801a/src/context_implementations.jl#L247
• In observe, where we increment it: https://github.com/TuringLang/DynamicPPL.jl/blob/24a73809b1a9f0f5b4e5f907f737b57c6eaf801a/src/context_implementations.jl#L272

In other words, we also need num_produce as a counter, so that we can assign the correct orders every time we add a variable.

Moreover, it's not clear to me if moving num_produce to AdvancedPS would be very useful unless we also move varinfo.orders. That could probably be done with some clever usage of contexts, but it's significantly more work.

There's also a base question of is the ordering of observe/assume statements something that the model/VarInfo should keep track of, or should it be the sampler's concern. I think the state of the sampler, so where num_produce is for a current particle, should be the sampler's concern. But I'm less clear on order. I think this goes somewhat deep into the question of what really is a Turing.jl model. If we do want to keep order in VarInfo then we have to think about how to do that when using VarNamedVector, which currently does not track order.

Thoughts?

[penelopeysm]

I think this is subsumed by https://github.com/TuringLang/DynamicPPL.jl/issues/744 which has a concrete proposal for what to do with num_produce