Using numpyro primitives in jax's control flows fori_loop, while_loop, and cond

Currently, we have scan implemented in contrib.control_flow. After #878, it should be doable to use other control flows in HMC/NUTS. The implementation should mimic what we currently have in scan (but will be much simpler) and can be addressed step-by-step:

• [ ] Support while_loop
  • We can assume no primitive in cond_fn because any primitive in cond_fn can be moved to body_fn
  • We can assume body_fn only contains observed nodes because while_loop is not meant to be used to collect stuffs.
  • Because body_fn only contains observed nodes, no need to worry about enumeration.
  • Because body_fn does not have latent variables, no need to worry about substitute/conditon stuffs.
• [ ] Support fori_loop
  • Any fori_loop can be rewrite as a while_loop. So supporting this should be straightforward.
• Support cond?
  • Support cond
  • Support cond with discrete latent variable inside body_fn
  • Support cond with discrete latent variable outside body_fn
  • Support cond under a plate of conditions
• [ ] Support nested while_loop
• Support nested scan with no latent variables (but we can collect stuff)

I think this is an interesting work and not as complicated as scan. This is a good chance to get familiar with JAX's control flow and many numpyro handlers. If anyone is interested in addressing this issue, please let me know any question that you have.

Start by reading contrib/control_flow/scan.py and the linked issue context, then compare its approach with JAX's while_loop, fori_loop, and cond control flows. Work through the checklist from while_loop and fori_loop to nested cases and cond cases with discrete latent variables; done means the listed control-flow scenarios are supported in HMC/NUTS.