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)