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

@junpenglao After checking your post and other post, now I am able to develop the efficient codes which works perfect for MCMC (thank you so much!!), now I want to use the same code for Variational Inference, but run into the following error message: Discrete variables are not supported by VI: count ~ Multinomial. Do you have any suggestions on what i shall do in this case? Just to remind you - I am attempting to implement supervised LDA using Variational Inference. Thanks so much!!!

def LDA_GLM(omega, y, K, M, N_V, n, alpha, gamma):
	with pm.Model() as model:
		eta = pm.Normal('b', mu=0, sigma=1, shape=K) # coefficient of linear regression on y
		sigma2 = pm.InverseGamma('sigma2',alpha =1.2, beta= 1.5) # variance of y 
		phi = pm.distributions.Dirichlet('phi', a=gamma, shape=(K, N_V)) # topic word matrix
		theta = pm.distributions.Dirichlet('theta', a=alpha, shape=(M, K)) # topic document matrix 
		omega = pm.DensityDist("doc", logp_lda_doc(phi, theta), observed=doc_t) # word document
        X_count = pm.Multinomial("count", n, p=theta, shape=(M, K)) # this line of code did not work for VI 
		X = pm.Deterministic("X", X_count/n)
		Y = pm.Normal('y', mu= pm.math.dot(X, eta), sigma = sigma2, shape =M, observed =y) # outcome variable 
	return model


ldaglm = LDA_GLM(omega2, Y, W, K, M, N_V,  n, alpha, gamma_null)
with ldaglm:
    ldaglm_vi = pm.fit(method="advi")