Dirichlet multidimensional tensor


Is there any way to give a Dirichlet prior to a tensor multidimensional form.

For example:

import numpy as np
import pymc3 as pm
U = 1000
B = 50
K = 10 # Number of assumed clusters
D = 3 
with pm.Model() as dir_model:
   theta_users = pm.Dirichlet('theta_u', a = np.ones((U,K,D)), shape = (U,K,D))
   theta_bts = pm.Dirichlet('theta_b', a = np.ones((B,K,D)), shape = (B,K,D))

Is this approach correct? With higher U, K, D will it be a problem?

I don’t think there should be a problem with that. Does it work?