def test_draw_samples(self, dtype, mean, mean_isSamples, var, var_isSamples, rv_shape, num_samples): n_dim = 1 + len(rv_shape) mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim) var_np = numpy_array_reshape(var, var_isSamples, n_dim) rand = np.random.randn(num_samples, *rv_shape) rv_samples_np = mean_np + rand * np.sqrt(var_np) rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype)) normal = Normal.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) var_mx = mx.nd.array(var, dtype=dtype) if not var_isSamples: var_mx = add_sample_dimension(mx.nd, var_mx) variables = {normal.mean.uuid: mean_mx, normal.variance.uuid: var_mx} rv_samples_rt = normal.draw_samples( F=mx.nd, variables=variables, num_samples=num_samples) assert np.issubdtype(rv_samples_rt.dtype, dtype) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_log_pdf(self, dtype, mean, mean_isSamples, var, var_isSamples, rv, rv_isSamples, num_samples): from scipy.stats import norm isSamples_any = any([mean_isSamples, var_isSamples, rv_isSamples]) rv_shape = rv.shape[1:] if rv_isSamples else rv.shape n_dim = 1 + len(rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape) mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim) var_np = numpy_array_reshape(var, var_isSamples, n_dim) rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim) log_pdf_np = norm.logpdf(rv_np, mean_np, np.sqrt(var_np)) normal = Normal.define_variable(shape=rv_shape, dtype=dtype).factor mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) var_mx = mx.nd.array(var, dtype=dtype) if not var_isSamples: var_mx = add_sample_dimension(mx.nd, var_mx) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_isSamples: rv_mx = add_sample_dimension(mx.nd, rv_mx) variables = {normal.mean.uuid: mean_mx, normal.variance.uuid: var_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) == isSamples_any if isSamples_any: assert get_num_samples(mx.nd, log_pdf_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples_with_broadcast(self, dtype, mean, mean_isSamples, var, var_isSamples, rv_shape, num_samples): mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) var_mx = mx.nd.array(var, dtype=dtype) if not var_isSamples: var_mx = add_sample_dimension(mx.nd, var_mx) var = var_mx.asnumpy() isSamples_any = any([mean_isSamples, var_isSamples]) 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(var), np.expand_dims(rand, axis=-1)).sum(-1) normal = MultivariateNormal.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = {normal.mean.uuid: mean_mx, normal.covariance.uuid: var_mx} draw_samples_rt = normal.draw_samples(F=mx.nd, variables=variables) assert np.issubdtype(draw_samples_rt.dtype, dtype) assert is_sampled_array(mx.nd, draw_samples_rt) == isSamples_any if isSamples_any: assert get_num_samples( mx.nd, draw_samples_rt) == num_samples, (get_num_samples( mx.nd, draw_samples_rt), num_samples)
def test_draw_samples_non_mock(self, plot=False): # Also make sure the non-mock sampler works dtype = np.float32 num_samples = 100000 mean = np.array([0.5]) variance = np.array([2]) rv_shape = (1,) mean_mx = add_sample_dimension(mx.nd, mx.nd.array(mean, dtype=dtype)) variance_mx = add_sample_dimension(mx.nd, mx.nd.array(variance, dtype=dtype)) rand_gen = None var = Normal.define_variable(shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor variables = {var.mean.uuid: mean_mx, var.variance.uuid: variance_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples assert rv_samples_rt.dtype == dtype from scipy.stats import norm if plot: plot_univariate(samples=rv_samples_rt, dist=norm, loc=mean[0], scale=np.sqrt(variance[0])) mean_est, scale_est = norm.fit(rv_samples_rt.asnumpy().ravel()) mean_tol = 1e-1 variance_tol = 1e-1 assert np.abs(mean[0] - mean_est) < mean_tol assert np.abs(variance[0] - scale_est ** 2) < variance_tol
def test_draw_samples_no_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() # n_dim = 1 + len(rv.shape) if is_samples_any else len(rv.shape) rand = np.random.rand(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype)) # rand_exp = np.expand_dims(rand, axis=-1) # lmat = np.linalg.cholesky(precision) # temp1 = np.matmul(lmat, rand_exp).sum(-1) # rv_samples_np = mean + temp1 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) assert get_num_samples(mx.nd, draw_samples_rt) == num_samples, \ (get_num_samples(mx.nd, draw_samples_rt), num_samples)
def test_draw_samples_non_mock(self, plot=False): # Also make sure the non-mock sampler works dtype = np.float32 num_samples = 100000 mean = np.array([0.5, 0]) covariance = np.array([[2, 0.5], [0.5, 2]]) rv_shape = (2,) mean_mx = add_sample_dimension(mx.nd, mx.nd.array(mean, dtype=dtype)) covariance_mx = add_sample_dimension(mx.nd, mx.nd.array(covariance, dtype=dtype)) rand_gen = None var = MultivariateNormal.define_variable(shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor variables = {var.mean.uuid: mean_mx, var.covariance.uuid: covariance_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples assert rv_samples_rt.dtype == dtype from scipy.stats import multivariate_normal if plot: plot_bivariate(samples=rv_samples_rt, dist=multivariate_normal, mean=mean, cov=covariance) # mean_est, scale_est = multivariate_normal.fit(rv_samples_rt.asnumpy().ravel()) mean_est = np.mean(rv_samples_rt.asnumpy(), axis=0) cov_est = np.cov(rv_samples_rt.asnumpy(), rowvar=False) mean_tol = 1e-1 covariance_tol = 1e-1 assert np.allclose(mean, mean_est, atol=mean_tol) assert np.allclose(covariance, cov_est, covariance_tol)
def test_draw_samples_non_mock(self, plot=False): # Also make sure the non-mock sampler works dtype = np.float32 num_samples = 100000 a = np.array([2]) b = np.array([5]) rv_shape = (1, ) a_mx = add_sample_dimension(mx.nd, mx.nd.array(a, dtype=dtype)) b_mx = add_sample_dimension(mx.nd, mx.nd.array(b, dtype=dtype)) rand_gen = None var = Beta.define_variable(shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor variables = {var.alpha.uuid: a_mx, var.beta.uuid: b_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples assert rv_samples_rt.dtype == dtype if plot: plot_univariate(samples=rv_samples_rt, dist=beta, a=a[0], b=b[0]) a_est, b_est, _, _ = beta.fit(rv_samples_rt.asnumpy().ravel()) a_tol = 0.2 b_tol = 0.2 assert np.abs(a[0] - a_est) < a_tol assert np.abs(b[0] - b_est) < b_tol
def test_log_pdf(self, dtype, low, low_is_samples, high, high_is_samples, rv, rv_is_samples, num_samples): is_samples_any = any([low_is_samples, high_is_samples, rv_is_samples]) rv_shape = rv.shape[1:] if rv_is_samples else rv.shape n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape) low_np = numpy_array_reshape(low, low_is_samples, n_dim) high_np = numpy_array_reshape(high, high_is_samples, n_dim) scale_np = high_np - low_np rv_np = numpy_array_reshape(rv, rv_is_samples, n_dim) # Note uniform.logpdf takes loc and scale, where loc=a and scale=b-a log_pdf_np = uniform.logpdf(rv_np, low_np, scale_np) var = Uniform.define_variable(shape=rv_shape, dtype=dtype).factor low_mx = mx.nd.array(low, dtype=dtype) if not low_is_samples: low_mx = add_sample_dimension(mx.nd, low_mx) high_mx = mx.nd.array(high, dtype=dtype) if not high_is_samples: high_mx = add_sample_dimension(mx.nd, high_mx) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_is_samples: rv_mx = add_sample_dimension(mx.nd, rv_mx) variables = {var.low.uuid: low_mx, var.high.uuid: high_mx, var.random_variable.uuid: rv_mx} log_pdf_rt = var.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 if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples(self): np.random.seed(0) samples_1_np = np.random.randn(5) samples_1 = mx.nd.array(samples_1_np) samples_2_np = np.random.randn(50) samples_2 = mx.nd.array(samples_2_np) m = Model() v = Variable(shape=(1,)) m.v2 = Normal.define_variable(mean=v, variance=mx.nd.array([1]), rand_gen=MockMXNetRandomGenerator(samples_1)) m.v3 = Normal.define_variable(mean=m.v2, variance=mx.nd.array([0.1]), shape=(10,), rand_gen=MockMXNetRandomGenerator(samples_2)) np.random.seed(0) v_np =np.random.rand(1) v_mx = mx.nd.array(v_np) v_rt = add_sample_dimension(mx.nd, v_mx) variance = m.v2.factor.variance variance2 = m.v3.factor.variance variance_rt = add_sample_dimension(mx.nd, variance.constant) variance2_rt = add_sample_dimension(mx.nd, variance2.constant) samples = m.draw_samples(F=mx.nd, num_samples=5, targets=[m.v3.uuid], variables={v.uuid: v_rt, variance.uuid: variance_rt, variance2.uuid: variance2_rt})[0] samples_np = v_np + samples_1_np[:, None] + np.sqrt(0.1)*samples_2_np.reshape(5,10) assert array_has_samples(mx.nd, samples) and get_num_samples(mx.nd, samples)==5 assert np.allclose(samples.asnumpy(), samples_np)
def test_compute_log_prob(self): m = Model() v = Variable(shape=(1,)) m.v2 = Normal.define_variable(mean=v, variance=mx.nd.array([1])) m.v3 = Normal.define_variable(mean=m.v2, variance=mx.nd.array([1]), shape=(10,)) np.random.seed(0) v_mx = mx.nd.array(np.random.randn(1)) v2_mx = mx.nd.array(np.random.randn(1)) v3_mx = mx.nd.array(np.random.randn(10)) v_rt = add_sample_dimension(mx.nd, v_mx) v2_rt = add_sample_dimension(mx.nd, v2_mx) v3_rt = add_sample_dimension(mx.nd, v3_mx) variance = m.v2.factor.variance variance2 = m.v3.factor.variance variance_rt = add_sample_dimension(mx.nd, variance.constant) variance2_rt = add_sample_dimension(mx.nd, variance2.constant) log_pdf = m.log_pdf(F=mx.nd, variables={m.v2.uuid: v2_rt, m.v3.uuid:v3_rt, variance.uuid: variance_rt, variance2.uuid: variance2_rt, v.uuid: v_rt}).asscalar() variables = {m.v2.factor.mean.uuid: v_rt, m.v2.factor.variance.uuid: variance_rt, m.v2.factor.random_variable.uuid: v2_rt} log_pdf_1 = mx.nd.sum(m.v2.factor.log_pdf(F=mx.nd, variables=variables)) variables = {m.v3.factor.mean.uuid: v2_rt, m.v3.factor.variance.uuid: variance2_rt, m.v3.factor.random_variable.uuid: v3_rt} log_pdf_2 = mx.nd.sum(m.v3.factor.log_pdf(F=mx.nd, variables=variables)) assert log_pdf == (log_pdf_1 + log_pdf_2).asscalar()
def test_log_pdf(self, dtype, prob_true, prob_true_is_samples, rv, rv_is_samples, num_samples): rv_shape = rv.shape[1:] if rv_is_samples else rv.shape n_dim = 1 + len(rv.shape) if not rv_is_samples else len(rv.shape) prob_true_np = numpy_array_reshape(prob_true, prob_true_is_samples, n_dim) rv_np = numpy_array_reshape(rv, rv_is_samples, n_dim) rv_full_shape = (num_samples, ) + rv_shape rv_np = np.broadcast_to(rv_np, rv_full_shape) log_pdf_np = bernoulli.logpmf(k=rv_np, p=prob_true_np) var = Bernoulli.define_variable(0, shape=rv_shape, dtype=dtype).factor prob_true_mx = mx.nd.array(prob_true, dtype=dtype) if not prob_true_is_samples: prob_true_mx = add_sample_dimension(mx.nd, prob_true_mx) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_is_samples: rv_mx = add_sample_dimension(mx.nd, rv_mx) variables = { var.prob_true.uuid: prob_true_mx, var.random_variable.uuid: rv_mx } log_pdf_rt = var.log_pdf(F=mx.nd, variables=variables) assert np.issubdtype(log_pdf_rt.dtype, dtype) assert get_num_samples(mx.nd, log_pdf_rt) == num_samples assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy())
def test_log_pdf_no_broadcast(self, dtype, mean, mean_isSamples, var, var_isSamples, rv, rv_isSamples, num_samples): mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) mean = mean_mx.asnumpy() var_mx = mx.nd.array(var, dtype=dtype) if not var_isSamples: var_mx = add_sample_dimension(mx.nd, var_mx) var = var_mx.asnumpy() rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_isSamples: rv_mx = add_sample_dimension(mx.nd, rv_mx) rv = rv_mx.asnumpy() from scipy.stats import multivariate_normal isSamples_any = any([mean_isSamples, var_isSamples, rv_isSamples]) rv_shape = rv.shape[1:] n_dim = 1 + len( rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape) mean_np = numpy_array_reshape(mean, isSamples_any, n_dim) var_np = numpy_array_reshape(var, isSamples_any, n_dim) rv_np = numpy_array_reshape(rv, isSamples_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], var_np[s][i])) r.append(a) log_pdf_np = np.array(r) normal = MultivariateNormal.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = { normal.mean.uuid: mean_mx, normal.covariance.uuid: var_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 is_sampled_array(mx.nd, log_pdf_rt) == isSamples_any if isSamples_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_isSamples, var, var_isSamples, rv_shape, num_samples): mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) var_mx = mx.nd.array(var, dtype=dtype) if not var_isSamples: var_mx = add_sample_dimension(mx.nd, var_mx) var = var_mx.asnumpy() rand = np.random.rand(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand.flatten(), dtype=dtype)) normal = MultivariateNormal.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = {normal.mean.uuid: mean_mx, normal.covariance.uuid: var_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 is_sampled_array(mx.nd, draw_samples_rt) == True
def test_draw_samples(self, dtype, mean, mean_is_samples, precision, precision_is_samples, rv_shape, num_samples): n_dim = 1 + len(rv_shape) mean_np = numpy_array_reshape(mean, mean_is_samples, n_dim) precision_np = numpy_array_reshape(precision, precision_is_samples, n_dim) rand = np.random.randn(num_samples, *rv_shape) rv_samples_np = mean_np + rand * np.power(precision_np, -0.5) rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype)) var = NormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor 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) variables = {var.mean.uuid: mean_mx, var.precision.uuid: precision_mx} rv_samples_rt = var.draw_samples( F=mx.nd, variables=variables, num_samples=num_samples) assert np.issubdtype(rv_samples_rt.dtype, dtype) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_log_pdf(self, dtype, mean, mean_is_samples, precision, precision_is_samples, rv, rv_is_samples, num_samples): is_samples_any = any([mean_is_samples, precision_is_samples, rv_is_samples]) rv_shape = rv.shape[1:] if rv_is_samples else rv.shape n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape) mean_np = numpy_array_reshape(mean, mean_is_samples, n_dim) precision_np = numpy_array_reshape(precision, precision_is_samples, n_dim) rv_np = numpy_array_reshape(rv, rv_is_samples, n_dim) log_pdf_np = norm.logpdf(rv_np, mean_np, np.power(precision_np, -0.5)) var = NormalMeanPrecision.define_variable(shape=rv_shape, dtype=dtype).factor 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) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_is_samples: rv_mx = add_sample_dimension(mx.nd, rv_mx) variables = {var.mean.uuid: mean_mx, var.precision.uuid: precision_mx, var.random_variable.uuid: rv_mx} log_pdf_rt = var.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 if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_eval(self, dtype, A, A_isSamples, B, B_isSamples, num_samples, broadcastable): np_isSamples = A_isSamples or B_isSamples if np_isSamples: if not A_isSamples: A_np = np.expand_dims(A, axis=0) else: A_np = A if not B_isSamples: B_np = np.expand_dims(B, axis=0) else: B_np = B res_np = np.einsum('ijk, ikh -> ijh', A_np, B_np) else: res_np = A.dot(B) eval = self._make_gluon_function_evaluation(dtype, broadcastable) A_mx = mx.nd.array(A, dtype=dtype) if not A_isSamples: A_mx = add_sample_dimension(mx.nd, A_mx) B_mx = mx.nd.array(B, dtype=dtype) if not B_isSamples: B_mx = add_sample_dimension(mx.nd, B_mx) variables = {eval.dot_input_0.uuid: A_mx, eval.dot_input_1.uuid: B_mx} res_rt = eval.eval(F=mx.nd, variables=variables) assert np_isSamples == is_sampled_array(mx.nd, res_rt) assert np.allclose(res_np, res_rt.asnumpy())
def test_draw_samples(self, dtype, low, low_is_samples, high, high_is_samples, rv_shape, num_samples): n_dim = 1 + len(rv_shape) low_np = numpy_array_reshape(low, low_is_samples, n_dim) high_np = numpy_array_reshape(high, high_is_samples, n_dim) rv_samples_np = np.random.uniform(low=low_np, high=high_np, size=(num_samples,) + rv_shape) rand_gen = MockMXNetRandomGenerator(mx.nd.array(rv_samples_np.flatten(), dtype=dtype)) var = Uniform.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor low_mx = mx.nd.array(low, dtype=dtype) if not low_is_samples: low_mx = add_sample_dimension(mx.nd, low_mx) high_mx = mx.nd.array(high, dtype=dtype) if not high_is_samples: high_mx = add_sample_dimension(mx.nd, high_mx) variables = {var.low.uuid: low_mx, var.high.uuid: high_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert np.issubdtype(rv_samples_rt.dtype, dtype) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_log_pdf_mean_variance(self, dtype, mean, mean_isSamples, variance, variance_isSamples, rv, rv_isSamples, num_samples): import scipy as sp isSamples_any = any([mean_isSamples, variance_isSamples, rv_isSamples]) rv_shape = rv.shape[1:] if rv_isSamples else rv.shape n_dim = 1 + len( rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape) mean_np = numpy_array_reshape(mean, mean_isSamples, n_dim) variance_np = numpy_array_reshape(variance, variance_isSamples, n_dim) rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim) beta_np = mean_np / variance_np alpha_np = mean_np * beta_np log_pdf_np = sp.stats.gamma.logpdf(rv_np, a=alpha_np, loc=0, scale=1. / beta_np) mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) variance_mx = mx.nd.array(variance, dtype=dtype) if not variance_isSamples: variance_mx = add_sample_dimension(mx.nd, variance_mx) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_isSamples: rv_mx = add_sample_dimension(mx.nd, rv_mx) gamma = GammaMeanVariance.define_variable(mean=mean_mx, variance=variance_mx, shape=rv_shape, dtype=dtype).factor variables = { gamma.mean.uuid: mean_mx, gamma.variance.uuid: variance_mx, gamma.random_variable.uuid: rv_mx } log_pdf_rt = gamma.log_pdf(F=mx.nd, variables=variables) assert np.issubdtype(log_pdf_rt.dtype, dtype) assert is_sampled_array(mx.nd, log_pdf_rt) == isSamples_any if isSamples_any: assert get_num_samples(mx.nd, log_pdf_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples_mean_variance(self, dtype, mean, mean_isSamples, variance, variance_isSamples, rv_shape, num_samples): n_dim = 1 + len(rv_shape) out_shape = (num_samples, ) + rv_shape mean_np = mx.nd.array(np.broadcast_to(numpy_array_reshape( mean, mean_isSamples, n_dim), shape=out_shape), dtype=dtype) variance_np = mx.nd.array(np.broadcast_to(numpy_array_reshape( variance, variance_isSamples, n_dim), shape=out_shape), dtype=dtype) gamma = GammaMeanVariance.define_variable(shape=rv_shape, dtype=dtype).factor mean_mx = mx.nd.array(mean, dtype=dtype) if not mean_isSamples: mean_mx = add_sample_dimension(mx.nd, mean_mx) variance_mx = mx.nd.array(variance, dtype=dtype) if not variance_isSamples: variance_mx = add_sample_dimension(mx.nd, variance_mx) variables = { gamma.mean.uuid: mean_mx, gamma.variance.uuid: variance_mx } mx.random.seed(0) rv_samples_rt = gamma.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) mx.random.seed(0) beta_np = mean_np / variance_np alpha_np = mean_np * beta_np rv_samples_mx = mx.nd.random.gamma(alpha=alpha_np, beta=beta_np, dtype=dtype) assert np.issubdtype(rv_samples_rt.dtype, dtype) assert is_sampled_array(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(rv_samples_mx.asnumpy(), rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_log_pdf(self, dtype_dof, dtype, degrees_of_freedom, random_state, scale_is_samples, rv_is_samples, num_data_points, num_samples, broadcast): # Create positive semi-definite matrices rv = make_spd_matrices_4d(num_samples, num_data_points, degrees_of_freedom, random_state=random_state) if broadcast: scale = make_spd_matrix(n_dim=degrees_of_freedom, random_state=random_state) else: scale = make_spd_matrices_4d(num_samples, num_data_points, degrees_of_freedom, random_state=random_state) degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom], dtype=dtype_dof) degrees_of_freedom = degrees_of_freedom_mx.asnumpy()[0] # ensures the correct dtype scale_mx = mx.nd.array(scale, dtype=dtype) if not scale_is_samples: scale_mx = add_sample_dimension(mx.nd, scale_mx) scale = scale_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 = scale_is_samples or rv_is_samples if broadcast: scale_np = np.broadcast_to(scale, rv.shape) else: n_dim = 1 + len(rv.shape) if is_samples_any and not rv_is_samples else len(rv.shape) scale_np = numpy_array_reshape(scale, is_samples_any, n_dim) rv_np = numpy_array_reshape(rv, is_samples_any, degrees_of_freedom) r = [] for s in range(num_samples): a = [] for i in range(num_data_points): a.append(wishart.logpdf(rv_np[s][i], df=degrees_of_freedom, scale=scale_np[s][i])) r.append(a) log_pdf_np = np.array(r) var = Wishart.define_variable(shape=rv.shape[1:], dtype=dtype, rand_gen=None).factor variables = {var.degrees_of_freedom.uuid: degrees_of_freedom_mx, var.scale.uuid: scale_mx, var.random_variable.uuid: rv_mx} log_pdf_rt = var.log_pdf(F=mx.nd, variables=variables) assert np.issubdtype(log_pdf_rt.dtype, dtype) assert is_sampled_array(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(self, dtype, a_shape, a_is_samples, b_shape, b_is_samples, rv_shape, num_samples): # Note: Tests above have been commented as they are very slow to run. # Note: Moved random number generation to here as the seed wasn't set if used above a = np.random.uniform(0.5, 2, size=a_shape) b = np.random.uniform(0.5, 2, size=b_shape) n_dim = 1 + len(rv_shape) a_np = numpy_array_reshape(a, a_is_samples, n_dim) b_np = numpy_array_reshape(b, b_is_samples, n_dim) rv_samples_np = np.random.beta(a_np, b_np, size=(num_samples, ) + rv_shape) var = Beta.define_variable(shape=rv_shape, dtype=dtype, rand_gen=None).factor a_mx = mx.nd.array(a, dtype=dtype) if not a_is_samples: a_mx = add_sample_dimension(mx.nd, a_mx) b_mx = mx.nd.array(b, dtype=dtype) if not b_is_samples: b_mx = add_sample_dimension(mx.nd, b_mx) variables = {var.a.uuid: a_mx, var.b.uuid: b_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert np.issubdtype(rv_samples_rt.dtype, dtype) assert is_sampled_array(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples rtol, atol = 1e-1, 1e-1 from itertools import product fits_np = [ beta.fit(rv_samples_np[:, i, j])[0:2] for i, j in (product(*map(range, rv_shape))) ] fits_rt = [ beta.fit(rv_samples_rt.asnumpy()[:, i, j])[0:2] for i, j in (product(*map(range, rv_shape))) ] assert np.allclose(fits_np, fits_rt, rtol=rtol, atol=atol)
def test_draw_samples_with_broadcast(self, dtype_dof, dtype, degrees_of_freedom, scale, scale_is_samples, rv_shape, num_samples): degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom], dtype=dtype_dof) scale_mx = mx.nd.array(scale, dtype=dtype) if not scale_is_samples: scale_mx = add_sample_dimension(mx.nd, scale_mx) rand = np.random.rand(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator(mx.nd.array(rand.flatten(), dtype=dtype)) reps = 1000 mins = np.zeros(reps) maxs = np.zeros(reps) for i in range(reps): rvs = wishart.rvs(df=degrees_of_freedom, scale=scale, size=num_samples) mins[i] = rvs.min() maxs[i] = rvs.max() # rv_samples_np = wishart.rvs(df=degrees_of_freedom, scale=scale, size=num_samples) var = Wishart.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = {var.degrees_of_freedom.uuid: degrees_of_freedom_mx, var.scale.uuid: scale_mx} draw_samples_rt = var.draw_samples(F=mx.nd, variables=variables) assert np.issubdtype(draw_samples_rt.dtype, dtype) assert is_sampled_array(mx.nd, draw_samples_rt) == scale_is_samples if scale_is_samples: assert get_num_samples(mx.nd, draw_samples_rt) == num_samples, (get_num_samples(mx.nd, draw_samples_rt), num_samples) assert mins.min() < draw_samples_rt.asnumpy().min() assert maxs.max() > draw_samples_rt.asnumpy().max()
def test_draw_samples(self, dtype, log_prob, log_prob_isSamples, rv_shape, num_samples, one_hot_encoding, normalization): n_dim = 1 + len(rv_shape) log_prob_np = numpy_array_reshape(log_prob, log_prob_isSamples, n_dim) rv_full_shape = (num_samples, ) + rv_shape log_prob_np = np.broadcast_to(log_prob_np, rv_full_shape[:-1] + (3, )) rand_np = np.random.randint(0, 3, size=rv_full_shape[:-1]) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand_np.flatten(), dtype=dtype)) if one_hot_encoding: rand_np = np.identity(3)[rand_np].reshape(*rv_full_shape) else: rand_np = np.expand_dims(rand_np, axis=-1) rv_samples_np = rand_np cat = Categorical.define_variable(0, num_classes=3, one_hot_encoding=one_hot_encoding, normalization=normalization, shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor log_prob_mx = mx.nd.array(log_prob, dtype=dtype) if not log_prob_isSamples: log_prob_mx = add_sample_dimension(mx.nd, log_prob_mx) variables = {cat.log_prob.uuid: log_prob_mx} rv_samples_rt = cat.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples assert np.allclose(rv_samples_np, rv_samples_rt.asnumpy())
def test_draw_samples_no_broadcast(self, dtype_dof, dtype, degrees_of_freedom, scale, scale_is_samples, rv_shape, num_samples): degrees_of_freedom_mx = mx.nd.array([degrees_of_freedom], dtype=dtype_dof) scale_mx = mx.nd.array(scale, dtype=dtype) if not scale_is_samples: scale_mx = add_sample_dimension(mx.nd, scale_mx) rand = np.random.rand(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand.flatten(), dtype=dtype)) var = Wishart.define_variable(shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = { var.degrees_of_freedom.uuid: degrees_of_freedom_mx, var.scale.uuid: scale_mx } draw_samples_rt = var.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) 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_eval_gluon_parameters(self, dtype, A, A_isSamples, B, B_isSamples, C, C_isSamples, num_samples, broadcastable): np_isSamples = A_isSamples or B_isSamples if np_isSamples: if not A_isSamples: A_np = np.expand_dims(A, axis=0) else: A_np = A if not B_isSamples: B_np = np.expand_dims(B, axis=0) else: B_np = B res_np = np.einsum('ijk, ikh -> ijh', A_np, B_np) if C_isSamples: res_np += C[:, :, None] else: res_np += C else: res_np = A.dot(B) if C_isSamples: res_np = res_np[None, :, :] + C[:, :, None] else: res_np += C np_isSamples = np_isSamples or C_isSamples eval = self._make_gluon_function_evaluation_rand_param( dtype, broadcastable) A_mx = mx.nd.array(A, dtype=dtype) if not A_isSamples: A_mx = add_sample_dimension(mx.nd, A_mx) B_mx = mx.nd.array(B, dtype=dtype) if not B_isSamples: B_mx = add_sample_dimension(mx.nd, B_mx) C_mx = mx.nd.array(C, dtype=dtype) if not C_isSamples: C_mx = add_sample_dimension(mx.nd, C_mx) variables = { eval.dot_input_0.uuid: A_mx, eval.dot_input_1.uuid: B_mx, eval.dot_const.uuid: C_mx } res_rt = eval.eval(F=mx.nd, variables=variables) assert np_isSamples == array_has_samples(mx.nd, res_rt) assert np.allclose(res_np, res_rt.asnumpy())
def test_log_pdf(self, dtype, log_prob, log_prob_isSamples, rv, rv_isSamples, num_samples, one_hot_encoding, normalization): rv_shape = rv.shape[1:] if rv_isSamples else rv.shape n_dim = 1 + len(rv.shape) if not rv_isSamples else len(rv.shape) log_prob_np = numpy_array_reshape(log_prob, log_prob_isSamples, n_dim) rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim) rv_full_shape = (num_samples, ) + rv_shape rv_np = np.broadcast_to(rv_np, rv_full_shape) log_prob_np = np.broadcast_to(log_prob_np, rv_full_shape[:-1] + (3, )) if normalization: log_pdf_np = np.log( np.exp(log_prob_np) / np.exp(log_prob_np).sum(-1, keepdims=True)).reshape(-1, 3) else: log_pdf_np = log_prob_np.reshape(-1, 3) if one_hot_encoding: log_pdf_np = (rv_np.reshape(-1, 3) * log_pdf_np).sum(-1).reshape( rv_np.shape[:-1]) else: bool_idx = np.arange(3)[None, :] == rv_np.reshape(-1, 1) log_pdf_np = log_pdf_np[bool_idx].reshape(rv_np.shape[:-1]) cat = Categorical.define_variable(0, num_classes=3, one_hot_encoding=one_hot_encoding, normalization=normalization, shape=rv_shape, dtype=dtype).factor log_prob_mx = mx.nd.array(log_prob, dtype=dtype) if not log_prob_isSamples: log_prob_mx = add_sample_dimension(mx.nd, log_prob_mx) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_isSamples: rv_mx = add_sample_dimension(mx.nd, rv_mx) variables = { cat.log_prob.uuid: log_prob_mx, cat.random_variable.uuid: rv_mx } log_pdf_rt = cat.log_pdf(F=mx.nd, variables=variables) assert np.issubdtype(log_pdf_rt.dtype, dtype) assert get_num_samples(mx.nd, log_pdf_rt) == num_samples assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy())
def test_log_pdf(self, dtype, alpha, alpha_isSamples, beta, beta_isSamples, rv, rv_isSamples, num_samples): import scipy as sp isSamples_any = any([alpha_isSamples, beta_isSamples, rv_isSamples]) rv_shape = rv.shape[1:] if rv_isSamples else rv.shape n_dim = 1 + len( rv.shape) if isSamples_any and not rv_isSamples else len(rv.shape) alpha_np = numpy_array_reshape(alpha, alpha_isSamples, n_dim) beta_np = numpy_array_reshape(beta, beta_isSamples, n_dim) rv_np = numpy_array_reshape(rv, rv_isSamples, n_dim) log_pdf_np = sp.stats.gamma.logpdf(rv_np, a=alpha_np, loc=0, scale=1. / beta_np) gamma = Gamma.define_variable(shape=rv_shape, dtype=dtype).factor alpha_mx = mx.nd.array(alpha, dtype=dtype) if not alpha_isSamples: alpha_mx = add_sample_dimension(mx.nd, alpha_mx) beta_mx = mx.nd.array(beta, dtype=dtype) if not beta_isSamples: beta_mx = add_sample_dimension(mx.nd, beta_mx) rv_mx = mx.nd.array(rv, dtype=dtype) if not rv_isSamples: rv_mx = add_sample_dimension(mx.nd, rv_mx) variables = { gamma.alpha.uuid: alpha_mx, gamma.beta.uuid: beta_mx, gamma.random_variable.uuid: rv_mx } log_pdf_rt = gamma.log_pdf(F=mx.nd, variables=variables) assert np.issubdtype(log_pdf_rt.dtype, dtype) assert is_sampled_array(mx.nd, log_pdf_rt) == isSamples_any if isSamples_any: assert get_num_samples(mx.nd, log_pdf_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(log_pdf_np, log_pdf_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples(self, dtype, alpha, alpha_isSamples, beta, beta_isSamples, rv_shape, num_samples): n_dim = 1 + len(rv_shape) out_shape = (num_samples, ) + rv_shape alpha_np = mx.nd.array(np.broadcast_to(numpy_array_reshape( alpha, alpha_isSamples, n_dim), shape=out_shape), dtype=dtype) beta_np = mx.nd.array(np.broadcast_to(numpy_array_reshape( beta, beta_isSamples, n_dim), shape=out_shape), dtype=dtype) gamma = Gamma.define_variable(shape=rv_shape, dtype=dtype).factor alpha_mx = mx.nd.array(alpha, dtype=dtype) if not alpha_isSamples: alpha_mx = add_sample_dimension(mx.nd, alpha_mx) beta_mx = mx.nd.array(beta, dtype=dtype) if not beta_isSamples: beta_mx = add_sample_dimension(mx.nd, beta_mx) variables = {gamma.alpha.uuid: alpha_mx, gamma.beta.uuid: beta_mx} mx.random.seed(0) rv_samples_rt = gamma.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) mx.random.seed(0) rv_samples_mx = mx.nd.random.gamma(alpha=alpha_np, beta=beta_np, dtype=dtype) assert np.issubdtype(rv_samples_rt.dtype, dtype) assert is_sampled_array(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples if np.issubdtype(dtype, np.float64): rtol, atol = 1e-7, 1e-10 else: rtol, atol = 1e-4, 1e-5 assert np.allclose(rv_samples_mx.asnumpy(), rv_samples_rt.asnumpy(), rtol=rtol, atol=atol)
def test_draw_samples(self, dtype, prob_true, prob_true_is_samples, rv_shape, num_samples): rv_full_shape = (num_samples, ) + rv_shape rand_np = np.random.normal(size=rv_full_shape) > 0 rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand_np.flatten(), dtype=dtype)) rv_samples_np = rand_np var = Bernoulli.define_variable(0, shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor prob_true_mx = mx.nd.array(prob_true, dtype=dtype) if not prob_true_is_samples: prob_true_mx = add_sample_dimension(mx.nd, prob_true_mx) variables = {var.prob_true.uuid: prob_true_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples assert np.array_equal(rv_samples_np, rv_samples_rt.asnumpy().astype(bool)) # Also make sure the non-mock sampler works rand_gen = None var = Bernoulli.define_variable(0, shape=rv_shape, rand_gen=rand_gen, dtype=dtype).factor prob_true_mx = mx.nd.array(prob_true, dtype=dtype) if not prob_true_is_samples: prob_true_mx = add_sample_dimension(mx.nd, prob_true_mx) variables = {var.prob_true.uuid: prob_true_mx} rv_samples_rt = var.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples) assert array_has_samples(mx.nd, rv_samples_rt) assert get_num_samples(mx.nd, rv_samples_rt) == num_samples assert rv_samples_rt.dtype == dtype