Strongly recommend you to try the marginalized HMM (see discussion in How to marginalized Hidden Markov Model with categorical?)
HMM model as in the main post use a discrete latent variable to represent the state, which cannot be sampled using HMC/NUTS - this usually means that you will get poor inference result. Marginalized HMM have other problems but a semi-supervised method should give reasonable inference result.