Behrens' Bayesian Learner Model : How to share parameters between steps?

I updated my notebook with a GaussianRW and it seems to do as well as the model with the betas (and sample much faster): Google Colaboratory

Neither model includes the volatility term, but so far I see no evidence that the integrated model won’t work (i.e., no need for fitting each step at a time, which would always discard information).