def test_with_mvnormal(self):
# 10 batch, 3-variate Gaussian
mu = np.random.randn(self.mixture_comps, 3)
mat = np.random.randn(3, 3)
cov = mat @ mat.T
chol = np.linalg.cholesky(cov)
w = np.ones(self.mixture_comps) / self.mixture_comps
with pm.Model() as model:
comp_dists = pm.MvNormal.dist(mu=mu, chol=chol, shape=(self.mixture_comps, 3))
mixture = pm.MixtureSameFamily(
"mixture", w=w, comp_dists=comp_dists, mixture_axis=0, shape=(3,)
)
prior = pm.sample_prior_predictive(samples=self.n_samples)
assert prior["mixture"].shape == (self.n_samples, 3)
assert mixture.random(size=self.size).shape == (self.size, 3)
if aesara.config.floatX == "float32":
rtol = 1e-4
else:
rtol = 1e-7
comp_logp = comp_dists.logp(model.test_point["mixture"].reshape(1, 3))
log_sum_exp = logsumexp(
comp_logp.eval() + np.log(w)[..., None], axis=0, keepdims=True
).sum()
assert_allclose(
model.logp(model.test_point),
log_sum_exp,
rtol,
)
def test_with_multinomial(self, batch_shape):
p = np.random.uniform(size=(*batch_shape, self.mixture_comps, 3))
n = 100 * np.ones((*batch_shape, 1))
w = np.ones(self.mixture_comps) / self.mixture_comps
mixture_axis = len(batch_shape)
with pm.Model() as model:
comp_dists = pm.Multinomial.dist(p=p, n=n, shape=(*batch_shape, self.mixture_comps, 3))
mixture = pm.MixtureSameFamily(
"mixture",
w=w,
comp_dists=comp_dists,
mixture_axis=mixture_axis,
shape=(*batch_shape, 3),
)
prior = pm.sample_prior_predictive(samples=self.n_samples)
assert prior["mixture"].shape == (self.n_samples, *batch_shape, 3)
assert mixture.random(size=self.size).shape == (self.size, *batch_shape, 3)
if aesara.config.floatX == "float32":
rtol = 1e-4
else:
rtol = 1e-7
comp_logp = comp_dists.logp(model.test_point["mixture"].reshape(*batch_shape, 1, 3))
log_sum_exp = logsumexp(
comp_logp.eval() + np.log(w)[..., None], axis=mixture_axis, keepdims=True
).sum()
assert_allclose(
model.logp(model.test_point),
log_sum_exp,
rtol,
)
def test_broadcasting_in_shape(self):
with pm.Model() as model:
mu = pm.Gamma("mu", 1.0, 1.0, shape=2)
comp_dists = pm.Poisson.dist(mu, shape=2)
mix = pm.MixtureSameFamily(
"mix", w=np.ones(2) / 2, comp_dists=comp_dists, shape=(1000,)
)
prior = pm.sample_prior_predictive(samples=self.n_samples)
assert prior["mix"].shape == (self.n_samples, 1000)
