pymc how do i add dims to gaussian mixture model

Question

I am struggling with the below issue. Suppose I have a gaussian mixture (3 components) but I want to do have individual weights per group (two groups in the below - so i would expect w to be of shape (2,3). How do I achieve this?

Tried with coords and with NormalMixture, and different variations of shape? Pymc 4.3 but also tried with pymc 5.x.

data=[45,56,78,99,45,64]
groups=[0,1,1,0,1,0]
with pm.Model() as model:
    mu = pm.Normal('mu', mu=50,sigma=10, shape=(3,))  
    components = pm.Normal.dist(mu=mu, shape=(3,))
    w = pm.Dirichlet('w', a=np.array([1, 1,1]))  # two mixture component weights.
    like = pm.Mixture('like', w=w, comp_dists=components, observed=data)
    trace=pm.sample()