I see, I didn’t know what you meant with augmented. For me that’s the normal Mixture model without marginalizing the indicator variable.
You can do that by sampling the latent variables and use those to index the parameters, a bit like the coal mining example: Introductory Overview of PyMC — PyMC 5.6.1 documentation
Otherwise if you want a mixture among different distributions and can’t just use a parameters switch, you can use CustomDist, something like
import pymc as pm
def dist(idx, mu1, mu2, sigma1, sigma2, size):
comp1 = pm.Normal.dist(mu1, sigma1, size=size)
comp2 = pm.Laplace(mu2, sigma2, size=size)
return pm.math.switch(idx, comp1, comp2)
with pm.Model() as m:
# Prior for idicator variables
idx = pm.Bernoulli("idx", p=0.5, size=(3,)) # can use categorical for more components
# Prior for mixture components parameters
mu1 = ...
mu2 = ...
sigma1= ...
sigma2 = ...
# Mixture likelihood with custom distribution
llike = pm.CustomDist("llike", idx, mu1, mu2, sigma1, sigma2, dist=dist, observed=[2, 2, 3])