How to speed up Pymc model compilation of a LDA topic model

It is however very much an issue for VI to handle discrete variable in practice as you most likely need gradient to train your model. Usually you need reparameterization trick like Gumbel-Softmax (see e.g., [1611.01144] Categorical Reparameterization with Gumbel-Softmax)

In your model, you sample a Multinomial with total count n, and then normalized it again by dividing n, which again gives you back the expectation of theta - modeling it this way is just extra stochastic-ness that get marginalized out, so I would advise to use theta directly and generate X_count after model inference.