Unnecessary computation in soft_vmap()

[tare]

Bug Description

I was confused why the map branch was executed

numpyro/numpyro/util.py

Line 453 in 627d19a

ys = lax.map(fn, xs) if num_chunks > 1 else fn(xs)

when I passed parallel=True to Predictive().

It looks like in that case num_chunks == 2 (should be 1) because pad == chunk_size (should be 0) even though batch_size % chunk_size == 0. This is the line that causes extra padding:

numpyro/numpyro/util.py

Line 441 in 627d19a

pad = chunk_size - (batch_size % chunk_size)

Please let me know whether my reasoning makes sense.

This commit fixes the issue.

Steps to Reproduce

N/A

Expected Behavior

No padding when batch_size % chunk_size == 0.