Log likelihood computation in numpyro can be extremely slow

[OriolAbril]

See #2367 for example. We should look into what makes it be slow and try to fix it, it would also be nice to add a progressbar so it is less surprising when things get stuck or freeze.

[aporsch1]

I think I'd like to work on this. I'll look at your guidelines and take a crack at it.

[OriolAbril]

Thanks! That would be greatly appreciated. I myself am not very familiar with numpyro but will do my best to help if I can

[aporsch1]

Great. We aren't parallelizing at all here, is there any downside to doing that? We could include batch size as a parameter or something like that which could lead to a decent size speedup, I think. I will also look more at numpyro to see if I can get more performance gains there.