Fix faulty interaction between jax.vmap and validate_args=True

[pierreglaser]

Closes #1684

@fehiepsi I changed the reproducer to a pattern that felt more prone to happen in practice than #1684 (comment):

    v_dist = jax.vmap(
        lambda loc, scale: dist.Normal(loc=loc, scale=scale, validate_args=True),
        in_axes=(0, 0),
    )(jnp.zeros((2,)), jnp.zeros((2,)))

[fehiepsi]

Nice fix! Thanks, @pierreglaser!