Fix a regression bug for ExpandedDistribution

[fehiepsi]

The bug is reported in forum

import numpyro.distributions as dist
from jax import random
import numpy as np

dist.Poisson(np.array([[1], [10], [100]])).expand([3, 10]).sample(random.PRNGKey(0))

returns

DeviceArray([[  2,  12,  88,   1,  12,  93,   0,  10, 113,   3],
             [  9, 117,   1,  10,  95,   2,   8,  88,   1,  11],
             [ 99,   1,   8,  95,   0,  12, 112,   1,   9, 105]],            dtype=int32)

This regression issue is introduced in #909 where some simplifications are performed and the tests do not cover the change.

[fritzo]

seems reasonable, but this might be cheaper using a single np.transpose rather than a loop of np.swapaxes calls.

[fehiepsi]

Thanks for reviewing and the helpful comment, @fritzo!