Specification of sampler

[reubenharry]

Current behavior

Sampling algorithms like hmc (and soon, mclmc) allow the user to specify an integrator. This is a little too much power, since if a user specified a non-euclidean integrator for hmc (or a euclidean one for mclmc), the algorithm would no longer be correct.

Desired behavior

It seems like what they should really be choosing are the coefficients of the symplectic integrator.

One way to do this would be to provide an enum IntegratorType = Enum('IntegratorType', ['mclachlan', 'yoshida', 'leapfrog'])

[AdrienCorenflos]

I disagree. The philosophy for the design of BlackJax has been to give the user as much flexibility as possible.
This is for at least two reasons:

• pedagogical purposes: if you want to demonstrate that HMC doesn't work with a non-euclidean integrator, then you need a framework that allows you to make the mistake on purpose.
• future integration of integrators (pun intended) is made harder by low-level nitpicking.

EDIT:

other reasons:

• longevity of the code: if you feed the arguments of the integrator to HMC rather than the integrator itself, you are saying that any change to the integrator API will need to result in a change to the HMC one.
• interface readability: passing the integrator reads much "cleaner" than passing a dictionary
• vmapping/jitting: it's hard to do on a subset of a pytree, much easier to do it via composition.

[AdrienCorenflos]

@junpenglao I am not sure where it's at but didn't you have design notes that arose from your work on vmapping NUTS at some point? All these points got raised then if I remember correctly, but I'd be damned if I can find a trace of this.

[junpenglao]

This is my note on unrolled NUTS but i am not sure if it is relevant to this discussion: https://github.com/tensorflow/probability/blob/main/discussion/technical_note_on_unrolled_nuts.md

[AdrienCorenflos]

I have in mind something specific to BlackJax, not tfp :D, well I guess it's lost for posterity

[reubenharry]

Thanks for your response! That makes sense.

I suppose if I have an objection, it's that to be really transparent, a (symplectic) integrator should be directly a function of the two updates (call them $O_1$ and $O_2$), whereas currently, it's a function of the log density of the target distribution. That means that the ODE that's being integrated really is different between the euclidean and non-euclidean integrators, which is why it just doesn't make sense to feed the non-euclidean integrator to hmc or vice-versa.

But I agree that passing functions directly, and having a maximally modular design is the goal.

[junpenglao]

Yeah i think the general approach we are taking here is that: the top level API for users just want to run the sampler, and if low level functions are used we assume the user knows what they are doing. @reubenharry I heard your concern tho but maybe let's improve it via documentation first?