Summing out latent discrete parameters in mixture model

Also here’s a notebook for a simpler toy model