Difficulty following quickstart

[twhentschel]

Hi all,

I just starting using blackjax a few weeks ago and it took me a while to understand the quickstart notebook - in particular, I barely understood the details in the inference loop until recently! Depending on who the target audience is, it might be worthwhile to add more comments so beginners to blackjax, jax, and HMC can understand the code in the quickstart a little better. For example, here is some roughly annotated code for the inference_loop

def inference_loop(rng_key, kernel, initial_state, num_samples):
    """
    Creates `num_samples` number of samples using an Markov Chain Monte Carlo (MCMC) `kernel`.
    
    rng_key: a random number generator key (random functions in jax require a rng key as the first argument,
            see https://jax.readthedocs.io/en/latest/jax.random.html for more details on pseudo-random number
            generation in jax)
    kernel: a MCMC kernel that generates a new state from its current state
    initial_state: the initial, starting state of the Markov Chain
    num_samples: number of states in the Markov Chain to generate

    returns: `num_samples` number of Markov Chain states 
    """
    @jax.jit
    def one_step(state, rng_key):
        """
        Make a single transition in the Markov Chain using a MCMC `kernel`, starting
        from `state`.

        state: current state of the Markov Chain
        rng_key: a random number generator key

        returns: a new state after a transition from `state` (return `state` twice because
        `jax.lax.scan()` expects a function that returns a result and a state to use
        for the next function call in the scan loop)
        """

        # a blackjax MCMC kernel returns the new state and information about the transition,
        # which we ignore
        state, _ = kernel(rng_key, state)
        return state, state

    # split(key, num) creates `num` RNG subkeys from `key`
    keys = jax.random.split(rng_key, num_samples)
    # scan over the RNG subkeys `keys` to generate MCMC states; `jax.lax.scan` carries along
    # the current state to perform each new transition
    _, states = jax.lax.scan(one_step, initial_state, keys)

    return states

This example might be overkill, but if there is some interest on incorporating more comments/guidance in the quickstart or other example notebooks I'd be happy to contribute!