NUTS is unacceptably slow – my model takes around 30 minutes to sample with NUTS, whereas I need it to run in under a minute for the application in question. It sounds like the JAX integration might help a lot, but it didn’t seem to speed up my model when I tried it (probably didn’t do it correctly).