Turns out it was easy to roll out everywhere, only took an hour or two. So now all done 
Note: for models where pi_semiz is dense (no zeros), this new approach is fractionally slower (say 1-3% slower), but I think that is worthwhile for the large speed gains when there are lots of zeros.