Fix padding calculation (#1993)

* Fix pad calculation

* Update paths to data files (#1996)

* Update paths to data files

* address numerical issues due to new jax release

* fix the size of the new SP500 file

* Add `sum_sites` option to sum loss over sites or return as dict. (#1995)

* Add `sum_sites` option to sum loss over sites or return as `dict`.

* Add missing newline in warning.

* Improve rng default in the `contrib/model.py` module (#1992)

* fix

* rm comment

* alternative rng

* suggestion

---------

Co-authored-by: Du Phan <[email protected]>
Co-authored-by: Till Hoffmann <[email protected]>
Co-authored-by: Juan Orduz <[email protected]>

Author: 4 people

numpyro/util.py

@@ -438,7 +438,7 @@ def soft_vmap(
     # it is better to catch OOM error and reduce chunk_size by half until OOM disappears.
     chunk_size = batch_size if chunk_size is None else min(batch_size, chunk_size)
     if chunk_size > 1:
-        pad = chunk_size - (batch_size % chunk_size)
+        pad = chunk_size - batch_size % chunk_size if batch_size % chunk_size else 0
         xs = jax.tree.map(
             lambda x: jnp.pad(x, ((0, pad),) + ((0, 0),) * (np.ndim(x) - 1)), xs
         )