Step Function for model with Categorical RV

NUTS is used alone because you don’t have any latent discrete variables as you were reasoning. As to why it is so painfully slow I don’t know. What’s the size of your dataset?

Edit: I see you specified shape. It doesn’t look like anything absurd.