# Shape issues when creating Dirichlet mixtures

I am having trouble writing models like this in PyMC3:

z_i\sim\text{Categorical}(p)
X_i|z_i\sim\text{Dirichlet}(\alpha_{z_i})

It’s a mixture model in which the components themselves are Dirichlet distributions. I studied this gaussian example from the documentation:

with pm.Model() as model:
mu = pm.Normal('mu', mu=np.arange(5), sigma=1, shape=5) # Each component has an independent mean
w = pm.Dirichlet('w', a=np.ones((3, 5)))  # w is a stack of 3 independent 5 component weight arrays
components = pm.Normal.dist(mu=mu, sigma=1, shape=(3, 5))
# The mixture is an array of 3 elements.
# Each can be thought of as an independent scalar mixture of 5 components
like = pm.Mixture('like', w=w, comp_dists = components, shape=3)


I tried to replicate it as follows:

with pm.Model() as model:
comp_dists = pm.Dirichlet.dist(np.ones((10,15)), shape=(10,15))
weights = pm.Dirichlet("weights", np.ones((20, 10)), shape=(20,10))
mix = pm.Mixture("likelihood", w=weights, comp_dists=comp_dists)


But it gives me ValueError: Input dimension mis-match. (input[0].shape[0] = 20, input[1].shape[0] = 10). I have tried exchanging the dimensions, using MixtureSameFamily instead of Mixture and using a list of dists instead os a multidimensional dist, but to no avail. Perhaps I am doing something silly with those shape arguments, but I’ve been stuck on this for a few days now.

Can you please provide some guidance? Thanks!