### Description @aseyboldt gave a nice explanation here that we could probably integrate in the docstrings: https://github.com/pyro-ppl/numpyro/pull/1751#issuecomment-1980569811