Just to be very specific the MCMC stage aims at estimating the distribution over parameters via samples from the posterior distribution:
Specify prior + likelihood
use the "model" the optimized solutions as the warmstarts of the K chains
warmstart the NUTS sampler with the window_adaption then sample using NUTS (blackjax)
these samples can then be used for prediction (note that for many samples and chains the memory requirements increase rapidly, so one should allow for compressed saving of those - maybe in chunks)
Different kind of "predictions" are possible. Let's maybe discuss this in a separate meeting or let's create a separate Issue for that.
Implement the following features: Given an input, a model and N sets of params, create an output distribution using MCMC.