Dirichlet with random shape

I have something that seemingly answers my previous question:

with pm.Model() as model:
	k = 5
	w = pm.Dirichlet('w', a=tt.stack([0.1,0.1,0.1,0.1,0.1]), shape=k)
	du = pm.DiscreteUniform('du',lower=0, upper=1,shape=k)
	idx_arr = []
	for i in range(k):
		if du[i]:
			idx_arr.append(i)

	mu = pm.Normal('mu', 0., 10., shape=k)
	tau = pm.Gamma('tau', 1., 1., shape=k)
	x_obs = pm.NormalMixture('x_obs', tt.stack([ w[h] for h in idx_arr ]), tt.stack([ mu[h] for h in idx_arr ]), tau=tt.stack([ tau[h] for h in idx_arr ]), observed=x)
	
with model:	
	trace = pm.sample(3000,n_init=10000, tune=500, random_seed=SEED)[1000:] 

With respect to the first response, does this implements the suggestion of “dynamically assign zeros”?