Easier Use of Distributions Under `vmap`

Distributions carry some static metadata (e.g., the batch shape), that make it a bit annoying to use in some other contexts. For example, they sometimes cause issues in dynestyx, where we use numpyro distributions quite generally.

An example of something I'd like to work, but doesn't, is just accessing variables of vmaped distributions:

def foo(mean):
    return dist.Normal(mean, jnp.ones_like(mean))

jax.vmap(foo)(jnp.zeros((2,2))).mean

Speaking offline with @fehiepsi, we agreed it would be desirable for such things to "just work." Looking back at previous discussions, however, it seems this is expected behavior under the current implementation (cf. https://github.com/pyro-ppl/numpyro/pull/1529#issuecomment-1411377292 or https://github.com/pyro-ppl/numpyro/issues/1684). Nevertheless, it would be nice to have a more extendable way of working with such things than imposing lots of numpyro-specific helpers on downstream users.

One way that might be possible is to try and defer shape evaluations when they may be inferred from parameter shapes (e.g., in all the scalar cases), in which case they become a lazy property. I had a quick vibe-coded go at this and it seems possible (see https://github.com/DanWaxman/numpyro/tree/dw-exploring-vmap-dists), though I haven't taken the time to go through that code very carefully (in particular, things that happen around expanded and transformed distributions), and some of the tests look wonky (in particular, I think Codex just changed test_distributions/test_promote_batch_shape_shares_data_and_preserves_input to pass ignoring all semantics haha), so please don't interpret it as a PR.

If that route seems interesting/acceptable, I'd be happy to take a closer/more careful look to see what it would take to do well and what it may break.

Start by reviewing the distribution shape-handling code and tests referenced in the issue, especially test_distributions/test_promote_batch_shape_shares_data_and_preserves_input, alongside the dw-exploring-vmap-dists branch. Compare the vmap example with behavior around expanded and transformed distributions; done means the example works without breaking the cited shape and data-preservation semantics and the relevant tests are reliable.