def test_mixture_inference() -> None:
mdo = MixtureDistributionOutput([GaussianOutput(), GaussianOutput()])
args_proj = mdo.get_args_proj()
args_proj.initialize()
args_proj.hybridize()
input = mx.nd.ones((BATCH_SIZE, 1))
distr_args = args_proj(input)
d = mdo.distribution(distr_args)
# plot_samples(d.sample())
trainer = mx.gluon.Trainer(
args_proj.collect_params(), "sgd", {"learning_rate": 0.02}
)
mixture_samples = mx.nd.array(np_samples)
N = 1000
t = tqdm(list(range(N)))
for i in t:
with mx.autograd.record():
distr_args = args_proj(input)
d = mdo.distribution(distr_args)
loss = d.loss(mixture_samples)
loss.backward()
loss_value = loss.mean().asnumpy()
t.set_postfix({"loss": loss_value})
trainer.step(BATCH_SIZE)
distr_args = args_proj(input)
d = mdo.distribution(distr_args)
obtained_hist = histogram(d.sample().asnumpy())
# uncomment to see histograms
# pl.plot(obtained_hist)
# pl.plot(EXPECTED_HIST)
# pl.show()
assert diff(obtained_hist, EXPECTED_HIST) < 0.5
def test_mixture_output(distribution_outputs) -> None:
mdo = MixtureDistributionOutput(*distribution_outputs)
args_proj = mdo.get_args_proj()
args_proj.initialize()
input = mx.nd.ones(shape=(512, 30))
distr_args = args_proj(input)
d = mdo.distribution(distr_args)
samples = d.sample(num_samples=NUM_SAMPLES)
sample = d.sample()
assert samples.shape == (NUM_SAMPLES, *sample.shape)
log_prob = d.log_prob(sample)
assert log_prob.shape == d.batch_shape