pyro-ppl / numpyro / 37214705711

93%

push

github

Skip the zero-width pad in soft_vmap (#2318)

* Skip the zero-width pad in soft_vmap

jnp.pad allocates and copies its input even when every pad width is
zero, so soft_vmap duplicated all of its inputs (e.g. the whole
posterior for Predictive(parallel=True)) whenever the batch size was a
multiple of the chunk size, which is always the case for the default
chunk_size=None. Only pad when there is something to pad.

* Convert soft_vmap inputs to jax arrays up front and test NumPy inputs

3 of 3 new or added lines in 1 file covered. (100.0%)

15654 of 16878 relevant lines covered (92.75%)

1.7 hits per line