Ah it looks like there is a working version of the high-dimensional GMMs in this issue: Mixture model with discrete data in Numpyro
I missed it in my search. The code in this thread is not accessing the correct dimensions created by the .plate() construct.