Palidrome diffusion solvers

[jeremiecoullon]

How to add diffusion solvers that are palidromes (ex: BAOAB or NOGIN)? These will usually (perhaps always?) need a gradient update halfway through the solver, so one solution is to have the update function in 2 parts. That way you calculate the gradient in between.

Example:

init_fn, update1, update2, get_params = baoab(1e-5)

for i in range(Nsamples):
  key, subkey1, subkey2 = random.split(key, 3)
  state = update1(i, subkey1, mygrad, state)
  mygrad = grad_log_post(get_params(state), *data) 
  state = update2(i, subkey2, mygrad, state)
  samples.append(get_params(state))

Some issues:

• sometimes (perhaps always?) the second update doesn't include randomness as it's after the "O" bit (which is the OU solver). So passing in a random key doesn't make sense.
• Should we write a second diffusion decorator (in diffusions.py), or extend the current one? Example: the decorator could take in an argument palidrome=False.
• Extending the diffusion decorator: if we do this, then this should also work nicely with the hmc diffusion, which currently returns resample_momentum as well as the usual 3 functions.

[jeremiecoullon]

This is merge in master. I created a "diffusion_factory" that generation diffusion decorators. This factory takes in arguments (is_palidrome and is_sghmc). The random key is passed in both updates (even though it seems that the second upate might never need randomness.