Use different sampler?

[swo]

          Also flagging for later that we may want to play with the choice of multinomial sampler here. In #50 I used numpyro's and it was somewhat ponderous for this many draws. There may not be much getting around that, but if we're going to do this many times on many datasets, it could be worth a few hour's testing between different library options and different batching schemes to save the down-stream computation.

Originally posted by @afmagee42 in #51 (comment)