First of all, thank you for this amazing library!

I have noticed a behavior that I cannot explain.

I’m trying to build a hierarchical Poisson regression and the PPC that I’m getting has the right shape, but the dimensions are mixed up (I can see different series mixed in within the same dimension).

Is this an expected behavior and should I always realign ppc myself?

Minimum example:

Let’s have 3 dimensional observed each with very different rate parameter (10,50,100)

```
# generate data
size=(10,)
rate_list=[10,50,100]
y_train=np.vstack([stats.poisson.rvs(rate,size=size) for rate in rate_list])
y_train.shape # gets you (3,10)
```

A simple model where we identify parameters separately for each series (independent thanks to the “series” plate)

```
# simple model
def model_poi(obs=None):
obs_dim=10
series_dim=3
with numpyro.plate('series',series_dim,dim=-2):
alpha=numpyro.sample('alpha',dist.Normal(0,3))
rate = numpyro.deterministic("rate",jnp.exp(alpha))
with numpyro.plate("time", obs_dim, dim=-1):
numpyro.sample("obs", dist.Poisson(rate), obs=obs)
```

Parameters are estimated with HMC and they are correct. PPC has the right shape as well

```
# get posterior predictive
predictive = Predictive(model_poi, posterior_samples,return_sites=["obs"],batch_ndims=1)
ppc = predictive(rng_key_, obs=None)
ppc['obs'].shape # gets you correct shape of (1000,3,10)
```

However, when you look at individual draws, you can see that the series are not represented row-wise anymore (middle dimension size = 3), you can see that each row goes series1,series2,series3, etc.

```
# show one draw
ppc['obs'][0]
DeviceArray([[ 6, 68, 110, 6, 43, 112, 2, 63, 99, 12],
[ 55, 100, 16, 58, 98, 13, 68, 97, 10, 57],
[107, 9, 56, 91, 10, 59, 85, 6, 54, 92]], dtype=int64)
```

My expectation would be that each series would be represented by a separate row like this:

DeviceArray([[ 6, 6, 2, 16, …],

[ 68, 43, 63, 55, …],

[110, 112, 99, 100, … ]], dtype=int64)

**My question:** Is this an expected behavior and should I always realign ppc myself?

Some additional observations:

- I can get the expected behavior by setting Predictive(… ,batch_ndims=2), which gives me the same shape but the dimensions are aligned correctly (as per the documentation, I’d have expected to get a shape like
`(num_chains x N x ...)`

but instead I get`(num_chains*N x ...)`

. This aligns with my expectation after inspecting the implementation, it’s just not what I would expect as per the documentation) - This issue goes away when I have some covariates which force the right shape on the latent variables that Posterior() can then read (in dim=-1)