See Sample size changes over NUTS MCMC sampling.
HMC/NUTS requires models with continuous parameters. If a model has discrete parameters, then HMC enumerates them in parallel to compute the log density of a model execution trace. This parallel enumeration requires designating and using additional batch dims to record subsequent parameter values in the model for each realization of the categorical variable.