Categorical distributional works strangely

I have faced this issue in a larger problem but for the sake of simplicity, I have prepared this small test. I basically don’t understand why categorical distribution gets affected by the mean of my MvNornal.
So the problem is this:
I sample from a categorical distribution and then choose to sample from a MvNormal conditioned on the categorical variable,

with pm.Model() as mm:
mu = theano.sharedt(mu)
z = pm.Categorical(‘z’, p=p)
d = pm.MvNormal(‘d’, mu=mu[z, :], cov=cov, shape=(1, 2))
trace = pm.sample(1000)

So I run the test with two different mu:
1- mu = np.array([[0, 0], [1, 1], [1, -1]])
2- mu = np.array([[0, 0], [10, 10], [10, -10]])
with the same Cov = cov = np.array([[1., 0.5], [0.5, 2]])

In the first one, the variable “z” is sampled correctly with the right proportions, in the second one it samples the “z” only from the category with the highest probability.

Please take a look at the following images


It works fine again if I change the cov = np.array([5, 0.5], [0.5, 5])

Unfortunately, mixing discrete RVs with continuous RVs can do weird stuff to HMC, and currently there is no good solution (other than marginalized out the discrete RVs):