The pymc3 version of that notebooks is here. Sampling categorical models like these tend to be difficult due to the categorical/discrete nature of some of the parameters. However, as you can see by comparing those 2 notebooks, there have been changes in v4. Whether they help speed-wise, I’m not sure. But it may be worth trying both.