Hi community,

I am have made a hierarchical version of a gaussian mixture model, and I would like to have access to the likelihood at each iteration of the MCMC, especially for when I run multiple chains, but I am getting errors due to broadcasting even though the model seems to work fine otherwise as it correctly identifies the mixture components on simulated data.

This is the model:

```
def Model_HGMM(K=2, dimension=2, data=None, label=None):
l = len(np.unique(label))
with numpyro.plate("components", K):
locs = numpyro.sample(
"locs",
dist.MultivariateNormal(jnp.zeros(dimension), 10 * jnp.eye(dimension)),
)
sigma = numpyro.sample("sigma", dist.LKJ(dimension, concentration=1))
with numpyro.plate("dimension", dimension):
with numpyro.plate("components", K):
with numpyro.plate("Age group", l):
locs_perturb = numpyro.sample(
"locs_perturb",
dist.Normal(loc=0,scale=0.1),
)
with numpyro.plate("Age group", l):
cluster_proba = numpyro.sample(
"cluster_proba", dist.Dirichlet(jnp.ones(K))
)
print("locs shape",locs.shape)
print("locs_perturb shape",locs_perturb.shape)
print("sigma shape",sigma.shape)
with numpyro.plate("data", len(data)):
assignment = numpyro.sample(
"assignment",
dist.Categorical(cluster_proba[label]),
infer={"enumerate": "parallel"},
)
print("locs[assignment] shape",locs[assignment].shape)
print("locs_perturb[label, assignment, :] shape",locs_perturb[label, assignment, :].shape)
print("sigma[assignment] shape",sigma[assignment].shape)
numpyro.sample(
"obs",
dist.MultivariateNormal(
locs[assignment] + locs_perturb[label, assignment, :],
covariance_matrix=sigma[assignment],
),
obs=data,
)
```

And this is the error I am getting when I am trying to get the likelihood using numpyro handlers.

```
ValueError: Incompatible shapes for broadcasting: shapes=[(600,), (600, 3)]
```

I have made a Jupyter notebook on simulated data to make it clearer.

Thanks in advance for your help!