MCLMC adaptation total num steps and initial guess

[hsimonfroy]

PR for #777

Fix MCLMC number of integration steps for tuning, and allow initial guess for parameters.

• We should be able to understand what the PR does from its title only;
• There is a high-level description of the changes;
• There are links to all the relevant issues, discussions and PRs;
• The branch is rebased on the latest main commit;
• Commit messages follow these guidelines;
• The code respects the current naming conventions;
• Docstrings follow the numpy style guide
• pre-commit is installed and configured on your machine, and you ran it before opening the PR;
• There are tests covering the changes;
• The doc is up-to-date;

I just don't know @reubenharry if you prefer make_L_step_size_adaptation to also return the number of steps as in adjusted_mclmc_make_L_step_size_adaptation, or to compute it outside the function. I opted for latter as the least modifying update.

[reubenharry]

Thanks very much for doing this! Looks good to me. My original design is a bit messy (the whole frac_tune thing), and maybe at some point I should clean it up, as well as aligning the API of adjusted_mclmc_make_L_step_size_adaptation better, but this seems good for now.

One question: why 2 effective samples?

[hsimonfroy]

One question: why 2 effective samples?

Not 2 effective samples, just 2 samples, because of the way ESS is computed:

blackjax/blackjax/diagnostics.py

Lines 160 to 162 in bd83572

	rho_hat = jnp.moveaxis(rho_hat, sample_axis, 0)
	rho_hat_even = rho_hat[0::2]
	rho_hat_odd = rho_hat[1::2]

I added an assert in effective_sample_size to handle this case, otherwise this raise an indexing error. Anyway, nobody should and would estimate ESS with 2 samples :')