I see… I think you can do that with a Potential or CustomDist. Something like this (untested):
def logp(H, mu, sigma, alpha):
logp_entries = pm.logp(pm.Beta.dist(mu=mu, sigma=sigma), H)
logp_rows = pm.logp(pm.Dirichlet.dist(alpha), H)
return logp_entries.sum(axis=-1) + logp_rows
with pm.Model() as m:
mu = ...
sigma = ...
alpha = ...
H = pm.CustomDist("H", mu, sigma, alpha, logp=logp, ndim_supp=1, ndims_params=[0, 0, 1])
That model adds the Dirichlet prior over rows (alpha.shape == H.shape=[-1]) if you had a single alpha just like a vanilla Dirichlet distribution with batched dimensions. I think the Stan example is doing the same, but please double check.
If you want the constraint over columns you can do that as well, just need to shape things around. By default the support dimensions (that is dimensions that are not independent) are to the rightmost position so this is more natural. Distribution Dimensionality — PyMC 5.0.1 documentation