Shape issues when creating Dirichlet mixtures

Questions
1

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!