Return adaptation extra information

[rlouf]

We currently only return the last state, the values of the parameter and the adapted kernel. However, the full chain and intermediate adaptation states can be useful when debugging inference.

In this PR we make window_adaptation, meads_adaptation and pathfinder_adaptation return this extra information. Closes #401.

The adaptation state and info are returned in the following NamedTuples:

class AdaptationResults(NamedTuple):
    state: PyTree
    kernel: Callable
    parameters: dict


class AdaptationInfo(NamedTuple):
    state: NamedTuple
    info: NamedTuple
    adaptation_state: NamedTuple

so the adaptation is ran e.g. as:

warmup = blackjax.window_adaptation(blackjax.nuts. logdensity_fn)
(state, kernel, parameters), info = warmup.run(rng_key, initial_position, 1000)

# OR
results, info = warmup.run(rng_key, initial_position, 1000)
state = results.state
kernel = results.kernel

IMO this feels clunky and a bit too high-level for Blackjax, mainly because of the fact that we return a kernel. I'm wondering if we should deprecate this way of doing things and favor instead the more explicit:

warmup = blackjax.window_adaptation(blackjax.nuts. logdensity_fn)
(state, parameters), info = warmup.run(rng_key, initial_position, 1000)
nuts = blackjax.nuts(logdensity_fn, **parameters)
hmc = blackjax.hmc(logdensity_fn, num_integration_steps=10, **parameters)  # we can also do this

[junpenglao]

+1 to deprecating returning a kernel. We had similar discussion some time back because returning kernel makes it difficult to vmap the warmup result (hence motivated the change #276).

I think it should be reasonable to do this instead:

warmup = blackjax.window_adaptation(blackjax.nuts, logdensity_fn)
(state, parameters), info = warmup.run(rng_key, initial_position, 1000)
kernel = warmup._kernel_fn(**parameters)

# The below is not garentee to work (only works if the kernel are within the same family like hmc/nuts)
nuts = blackjax.nuts(logdensity_fn, **parameters)
hmc = blackjax.hmc(logdensity_fn, num_integration_steps=10, **parameters)

WDYT?

[codecov]

Codecov Report

Merging #466 (5f615bc) into main (a4ee853) will decrease coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #466      +/-   ##
==========================================
- Coverage   99.16%   99.16%   -0.01%     
==========================================
  Files          48       48              
  Lines        1923     1918       -5     
==========================================
- Hits         1907     1902       -5     
  Misses         16       16              

Impacted Files	Coverage Δ
blackjax/__init__.py	100.00% <ø> (ø)
blackjax/adaptation/meads_adaptation.py	100.00% <ø> (ø)
blackjax/adaptation/__init__.py	100.00% <100.00%> (ø)
blackjax/kernels.py	99.18% <100.00%> (-0.02%)	⬇️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[junpenglao]

Basically, the abstraction being (mutation at different step):

# API level MCMC kernel
next_key, state = kernel.one_step(rng_key, last_state, **param)
# API level warmup kernel
next_param = tuning.one_step(param, state)

# User level sampling
tuned_param = blackjax.some_warmup_routine(...)
kernel_one_step = kernel.one_step(rng_key, last_state, **tuned_param)
for i in range(total_samples):
    rng_key, state = kernel_one_step(rng_key, last_state)

[junpenglao]

Here (and below) we are returning 2 tuple/namedtuple with similar structure - should we just return AdaptationInfo instead and do unpacking below?

[rlouf]

What are you suggesting here exactly? I'm not sure I understand.

[junpenglao]

I am suggesting to return

        return AdaptationInfo(new_state, info, new_adaptation_state)

[rlouf]

I like to only carry-over the values that are necessary in the scan loop. This often leads to repetitions but hasn't bothered me so far. What bothers you here?

[junpenglao]

oh right that's for using in scan - yeah than it is almost unavoidable

[rlouf]

+1 to deprecating returning a kernel. We had similar discussion some time back because returning kernel makes it difficult to vmap the warmup result (hence motivated the change #276).

I think it should be reasonable to do this instead:

kernel = warmup._kernel_fn(**parameters)

We might as well call blackjax.nuts directly here, warmup._kernel_fn is not doing any work.