I tried to implement toy models in emcee and numpyro as well. I also tried a version with the theano built-in argsort instead of the custom lexsort operation. It seems that the sorting part is what makes this model so slow. I am not sure yet how to solve this. Maybe there are more efficient ways to implement the median, otherwise maybe using GPUs with theano can help (suggested by Merwan).
I tried to implement toy models in emcee and numpyro as well. I also tried a version with the theano built-in
argsortinstead of the customlexsortoperation. It seems that the sorting part is what makes this model so slow. I am not sure yet how to solve this. Maybe there are more efficient ways to implement the median, otherwise maybe using GPUs with theano can help (suggested by Merwan).