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

I am trying to reproduce some results in Probabilistic-Programming-and-Bayesian-Methods-for-Hackers/Ch3_IntroMCMC_PyMC3.ipynb at master · CamDavidsonPilon/Probabilistic-Programming-and-Bayesian-Methods-for-Hackers · GitHub, 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