I don’t see where you marginalized the categorical variable. You would have to enumerate over all the possibilities and collect the respective conditional probabilities (or use a helper that does this like pm.Mixture), but you seem to have simply dropped that variable?