Feature Request: Allow user to specify batch size in 'vectorization' chain method

Feature Summary

When running MCMC with chain_method='vectorized', the likelihood function is automatically vmapped so that all chains are computed simultaneously. This feature would allow the user to specify how many chains are computed simultaneously.

Suggestion: Replace the following line
https://github.com/pyro-ppl/numpyro/blob/6b68d731823568c3988c91d877415779a774f7ce/numpyro/infer/mcmc.py#L416
with lax.map(fn, x, batch_size=self.vectorization_batch_size)

Why is this needed?

For complicated likelihood functions (e.g. in my application, my likelihood function involves a gravitational lensing image simulation pipeline starting with source/deflector parameters), vmapping over all chains can result in unnecessary memory usage and slowdowns if the GPU's computational power can be saturated with batch sizes much lower than the number of chains.

One may choose to use chain_method='sequential' instead, but in this case, the GPU's computational power may not be fully saturated running only one chain at a time, so a middle ground is desired. Furthermore, some MCMC methods do not support chain_method='sequential' (e.g. AIES).

Start at numpyro/infer/mcmc.py around line 416 and inspect how the vectorized chain method applies the likelihood across chains. Trace the MCMC configuration path to determine where a batch-size option belongs, then verify that vectorized execution uses the requested number of chains per batch without changing other chain methods.