Thanks! learning_weights, softmax_alpha, and word_order_prior all raise the NotImplementedError, while hyper_gammas doesn’t (or did you mean I should check something else?).
Unfortunately the version I posted on the gist is the stripped down version of the model - it is a pretty complicated model especially with the big scan operation. To be honest it’s been a bit of a nightmare to make this run on the real dataset (which is pretty large). I tried a bunch of things to make it run better, as I documented in this post. Eventually I just decided to go brute force and fit it just with pm.sample, it took about 1000 hours but convergence was very poor. This is despite extensive parameter recovery simulations (on smaller simulated datasets) working well.