No U-Turn Sampler

[dustinvtran]

Also see http://discuss.edwardlib.org/t/any-plans-for-nuts/.

References

• https://arxiv.org/abs/1111.4246

[korkinof]

I will eventually give it a try to implement this.
If anyone else wants to do this in the meantime, just give a shout here!

[ssequeira]

I'd be happy to try my hand at this.

[emilemathieu]

Hi !
I created a PR #728 with the implementation of the Dual Averaging algorithm for adapting the step size of HMC, which is a first step to implement NUTS.

[Razcle]

Hi I wondered if there was any update to this? I think there's probably a fair bit of demand for this feature and if help is required, I'd like to chip in.

[dustinvtran]

No need to put effort into. @axch has been working on this. Maybe he can reply on when something public-facing may appear in TF?

[axch]

Don't know about time lines, exactly, but work is proceeding.

[axch]

Actually, @dustinvtran, can you assign me this github issue (or give me sufficient permissions to assign it to myself)?

[dustinvtran]

Gave you permission. Looks like you need to accept before assignments are available.

FYI, if I recall correctly, @axch has mostly been working on the adaptive path length. For reference, it could be useful to look at @emilemathieu's implementation of the adaptive step size in #728.

[axch]

As an update, there is now a draft barebones NUTS implementation in TensorFlow Probability, at https://github.com/tensorflow/probability/tree/master/experimental/no_u_turn_sampler. This implementation is expected to be compatible with Edward2. See also the associated TFP issues tensorflow/probability#125, tensorflow/probability#126, tensorflow/probability#142, and tensorflow/probability#143.

[axch]

TFP now has a batch NUTS (path length adaptation, haven't tested step size adaptation yet but don't anticipate excessive problems) here: https://github.com/tensorflow/probability/blob/master/tensorflow_probability/python/experimental/mcmc/nuts.py.

[junpenglao]

NUTS sampler with multinominal sampling is currently on master: tensorflow/probability@e57eba6

You can have a look at this end2end test of how to do sampling with dualaveraging step size adaptation: https://github.com/tensorflow/probability/blob/e57eba6c0ff15b2c94e0d802115641683206c38a/tensorflow_probability/python/mcmc/nuts_test.py#L284-L381