Hierarchical Multinomial where groups change in each observation

This may be helpful for the shape issues: Distribution Dimensionality — PyMC 5.19.1 documentation