I was thinking of vectorizing and the problem is each n has a different value of t, I will have to keep the loop because of that. In the worst case, I will make each value of t the same. I realized that I can put any shape in the mean value of multivariate normal, so I change matrix normal to just multivariate normal and also changed discrete latent variables to continuous random variables and I also switched to ADVI.
Do you think if there is any another way to avoid instantiating n instances? I tried with n=400 and it is taking for instantiating, probably I would have to make all t values same.