Stan warmup

[rlouf]

In this Pull Request we implement the Stan warmup for a single chain. We also change the signature of the HMC and NUTS kernel factories. Example:

kernel_factory = lambda step_size, inverse_mass_matrix: nuts.kernel(potential_fn, step_size, inverse_mass_matrix)

# `info` is a tuple that contains the chain, all warmup states and the chain info
state, (step_size, inverse_mass_matrix), info = stan_warmup(
    rng_key,
    kernel_factory,
    initial_state,
    num_steps=1000
)

# We can use different kernels in the HMC family for warmup and sampling
kernel = hmc.kernel(potential_fn, steps_size, inverse_mass_matrix)

[codecov]

Codecov Report

Merging #16 (b064f27) into master (c6f75e9) will increase coverage by 1.31%.
The diff coverage is 98.36%.

@@            Coverage Diff             @@
##           master      #16      +/-   ##
==========================================
+ Coverage   95.47%   96.78%   +1.31%     
==========================================
  Files          10       15       +5     
  Lines         398      623     +225     
==========================================
+ Hits          380      603     +223     
- Misses         18       20       +2     

Impacted Files	Coverage Δ
blackjax/stan_warmup.py	96.19% <96.19%> (ø)
blackjax/adaptation/mass_matrix.py	100.00% <100.00%> (ø)
blackjax/adaptation/step_size.py	100.00% <100.00%> (ø)
blackjax/hmc.py	100.00% <100.00%> (+1.92%)	⬆️
blackjax/inference/proposal.py	100.00% <100.00%> (ø)
blackjax/nuts.py	100.00% <100.00%> (+1.85%)	⬆️
blackjax/optimizers/__init__.py	100.00% <100.00%> (ø)
blackjax/optimizers/dual_averaging.py	100.00% <100.00%> (ø)
... and 3 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update c6f75e9...b064f27. Read the comment docs.

[rlouf]

I need to separate "run" from the function that updates the stan warmup, in case we want to share information across chains.

[rlouf]

This is ready to review. We will just have to wait until #21 is merged to be able to test the Stan warmup with the NUTS kernel.