Welcome!
I might suggest looking through the Gaussian Mixture Model notebook, which has something very similar but written for version 4 of PyMC (and makes use of the now-default coords/dims rather than shapes). If that doesn’t help, let us know and someone can help get you what you need.