Stats from component samplers are lost with Gibbs

[penelopeysm]

Turing.jl/src/mcmc/gibbs.jl

Lines 323 to 335 in 37af7fd

	vi, states = gibbs_initialstep_recursive(
	rng,
	model,
	AbstractMCMC.step,
	varnames,
	samplers,
	vi;
	initial_params=initial_params,
	kwargs...,
	)
	transition = discard_sample ? nothing : DynamicPPL.ParamsWithStats(vi, model)
	return transition, GibbsState(vi, states)
	end

Here Gibbs returns a ParamsWithStats that only has parameters but no stats. In principle, it should aggregate stats from all component samplers, combine them into a single NamedTuple and return them.

This is incredibly annoying though because stats are provided via component samplers' transitions, not states, and #2759 makes sure that we don't actually use the parameters in component samplers' transitions, because that often requires reevaluating the model. Very, very annoying. The solution to this is to unify transitions and states, see TuringLang/AbstractMCMC.jl#149, but, mehhhhhh.

[penelopeysm]

As a stopgap measure, we could shove the same stats into the sampler state as well (i.e. duplicate them in both the transition and state). This should cost no extra model evaluations. Then Gibbs would have a function, Turing.Inference.gibbs_get_stats that extracts the stats from the state, which could be overloaded for other samplers.

This will mean that once again to make Gibbs work with external samplers, the external sampler will have to overload a Turing method, which is a bit disappointing. But assuming that transitions and states are unified some time in the future, this will probably be refactored into AbstractMCMC.getstats(state).