Follow-up issues for HMC about numpyro

[fehiepsi]

Alternatively (if this is relatively cheap), we can also just let the adapter function take in z and do the flattening.

I vote for this way. I think that it is cheap because we have a bunch of this operator in _is_turning function. It is easy to detect if it is cheap or not when we start to work with more complicated models.

The current hmc_kernel can have a separate get_kernel method that would return either an HMC or NUTS kernel depending on args. This can be called by the user or from inside mcmc.

Maybe we'd better discuss about this design choice in message. :)

Default kinetic energy function.

Yes, I agree.

Pyro style model/guide syntax

Yup.

[fehiepsi]

@neerajprad Look like most issues of HMC have been resolved. In upcoming days, I will address memory issue of iterative NUTS, and try to reduce its compiling time if possible. I think we have some room to improve given that functions will be recompiled when static args are not put in keywords (which we observed in dirichlet pr).

[neerajprad]

This is great, I am addressing the last remaining issue here.