Works like a charm!! Thanks a lot Ricardo.
One last question 
. How can I speed up this sampling? I know that discrete models cannot be sampled using numpyro for example, but it takes ages with hundreds/thousand of data points. Would I need to change it to some continuous model to benefit from faster sampling methods?