vmap produces wrong results silently

``` python
rng = jax.random.PRNGKey(0)
loc = jax.random.uniform(rng,(4,1))
logscale = jnp.array([1.])

pi1 = distrax.MultivariateNormalDiag(loc, jnp.exp(logscale))
def apply_dist(x,logscale):
return distrax.MultivariateNormalDiag(x, jnp.exp(logscale))
pi2 = jax.vmap(apply_dist, (0,None))(loc, jnp.exp(logscale))

print(pi1.sample(seed=rng).shape) #(4,1)
print(pi2.sample(seed=rng).shape) #(4,4)
```