What is equivalence to `pymc3.ElemwiseCategorical` in pyro?

I am trying to reproduce some results in https://github.com/CamDavidsonPilon/Probabilistic-Programming-and-Bayesian-Methods-for-Hackers/blob/master/Chapter3_MCMC/Ch3_IntroMCMC_PyMC3.ipynb, and they use pymc3.ElemwiseCategorical to sample categorical variables, just wonder whether there are similar things in Pyro, or we simply use NUTS?

Thank you!

there is a Categorical distribution. that is presumably what you want? see e.g. here for example usage