Seth Axen

Alternatively we could keep both the iterative and recursive NUTS implementations and only use the iterative one for vectorized NUTS (dispatching on matrix of positions). The advantage of this is...

Okay, NUTS is completely vectorized, and for a single chain, it produces the same output for recursive and iterative NUTS. I've tested with 100 parallel chains with mvnormal and mvcauchy...

Thanks @xukai92 for the review. I got busy and have had to put this on hold. I also think the PR might need to be broken into a few more...

> I am not familiar with Julia code, but probably you could avoid copying by doing [some](https://github.com/tensorflow/probability/blob/12954f1bae6b15f70d318fb9bed9f4788bf0cfad/tensorflow_probability/python/mcmc/nuts.py#L545-L549) [indexing](https://github.com/tensorflow/probability/blob/12954f1bae6b15f70d318fb9bed9f4788bf0cfad/tensorflow_probability/python/mcmc/nuts.py#L558-L561) and [update of tensor](https://github.com/tensorflow/probability/blob/12954f1bae6b15f70d318fb9bed9f4788bf0cfad/tensorflow_probability/python/mcmc/nuts.py#L662-L673). It's true, there may be other refactors that...

> After B is complete, you added that to A (`reverse(B)+A` or `A+B`, `+` is concatenation conceptually). So IIUC you are saying you will only do `A+B`, which I worry...

@junpenglao shared with me this related Stan discourse: https://discourse.mc-stan.org/t/parallel-dynamic-hmc-merits/10895/

I plan to implement this in three consecutive PRs, each with their own release. 1. Reimplement NUTS to be iterative. The implementation should be equivalent to the recursive implementation, so...

> Something like > > ```julia > df = (x = rand(100), y = rand(100), z = rand(100)) > vars = [:x, :y, :z] > data(df) * (visual(Scatter) + linear())...

> I think something like > > ```julia > data(df) * (visual(Scatter) + linear()) * mapping(:x, :x, col = :group => nonnumeric, row = :group => nonnumeric) > ``` >...

I'd like to contribute to this as I find the time. But before I get started, @MansMeg, how do you recommend proceeding? 1 model per PR, or are batches per...