Test using cached or pre-heated metric for NUTS sampling to decrease inference time

[SamuelBrand1]

From f2f discussion with @kokbent he noted that a large proportion of the sampling time in NUTS was during early sampling. This contributes to very long sampling times overall (many hours to get to 4000 samples).

This is suggestive that a lot of time is spent during adaptation phase, probably/possibly identification of a good metric for the momentum refresh step in NUTS.

A possible action to help with this would be to either serialise the metric from a successful run, for reusing in future runs on the same state. Or doing the same using an estimate of the inverse correlation matrix of posterior samples. A major limitation here would of course be that the number of parameters grows with number of variants, but this might be a useful approach to pre-heating the NUTS process.

[seabbs]

Just a fly by but a simple speed up is to reduce warmup steps as NUTs often doesn't need the default settings.

(otherwise yes agree with ideas here)

[kokbent]

Need inferer/runner to expose the inverse_mass_matrix argument. Will work on that.

[arik-shurygin]

Need inferer/runner to expose the inverse_mass_matrix argument. Will work on that.

Hey @kokbent is this functionality you found useful in the most recent update? If so I can expose the inverse_mass_matrix in main as well