def test_draw_samples_with_broadcast(self, dtype, mean, mean_is_samples,
precision, precision_is_samples,
rv_shape, num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
# precision = precision_mx.asnumpy()
is_samples_any = any([mean_is_samples, precision_is_samples])
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(
mx.nd.array(rand.flatten(), dtype=dtype))
# rv_samples_np = mean + np.matmul(np.linalg.cholesky(precision), np.expand_dims(rand, axis=-1)).sum(-1)
normal = MultivariateNormalMeanPrecision.define_variable(
shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {
normal.mean.uuid: mean_mx,
normal.precision.uuid: precision_mx
}
draw_samples_rt = normal.draw_samples(F=mx.nd, variables=variables)
assert np.issubdtype(draw_samples_rt.dtype, dtype)
assert array_has_samples(mx.nd, draw_samples_rt) == is_samples_any
if is_samples_any:
assert get_num_samples(mx.nd, draw_samples_rt) == num_samples, \
(get_num_samples(mx.nd, draw_samples_rt), num_samples)
def test_log_pdf_with_broadcast(self, dtype, mean, mean_is_samples, precision, precision_is_samples,
rv, rv_is_samples, num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
mean = mean_mx.asnumpy()
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
precision = precision_mx.asnumpy()
rv_mx = mx.nd.array(rv, dtype=dtype)
if not rv_is_samples:
rv_mx = add_sample_dimension(mx.nd, rv_mx)
rv = rv_mx.asnumpy()
is_samples_any = any([mean_is_samples, precision_is_samples, rv_is_samples])
rv_shape = rv.shape[1:]
n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape)
mean_np = np.broadcast_to(mean, (5, 3, 2))
precision_np = np.broadcast_to(precision, (5, 3, 2, 2))
rv_np = numpy_array_reshape(rv, is_samples_any, n_dim)
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype))
r = []
for s in range(len(rv_np)):
a = []
for i in range(len(rv_np[s])):
a.append(multivariate_normal.logpdf(rv_np[s][i], mean_np[s][i], np.linalg.inv(precision_np[s][i])))
r.append(a)
log_pdf_np = np.array(r)
normal = MultivariateNormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {
normal.mean.uuid: mean_mx, normal.precision.uuid: precision_mx, normal.random_variable.uuid: rv_mx}
log_pdf_rt = normal.log_pdf(F=mx.nd, variables=variables)
assert np.issubdtype(log_pdf_rt.dtype, dtype)
assert array_has_samples(mx.nd, log_pdf_rt) == is_samples_any
if is_samples_any:
assert get_num_samples(mx.nd, log_pdf_rt) == num_samples, (get_num_samples(mx.nd, log_pdf_rt), num_samples)
assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy())
def test_draw_samples_with_broadcast_no_numpy_verification(self, dtype, mean, mean_is_samples, precision,
precision_is_samples, rv_shape, num_samples):
mean_mx = mx.nd.array(mean, dtype=dtype)
if not mean_is_samples:
mean_mx = add_sample_dimension(mx.nd, mean_mx)
precision_mx = mx.nd.array(precision, dtype=dtype)
if not precision_is_samples:
precision_mx = add_sample_dimension(mx.nd, precision_mx)
# precision = precision_mx.asnumpy()
rand = np.random.rand(num_samples, *rv_shape)
rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype))
normal = MultivariateNormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor
variables = {normal.mean.uuid: mean_mx, normal.precision.uuid: precision_mx}
draw_samples_rt = normal.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples)
assert np.issubdtype(draw_samples_rt.dtype, dtype)
assert array_has_samples(mx.nd, draw_samples_rt)