def gpy_comb_kernel_test(X, X_isSamples, X2, X2_isSamples, kernel_params, num_samples, dtype, mf_kernel_create, gpy_kernel_create): X_mx = prepare_mxnet_array(X, X_isSamples, dtype) X2_mx = prepare_mxnet_array(X2, X2_isSamples, dtype) kern = mf_kernel_create().replicate_self() kernel_params_mx = {kern.name + '_' + k + '_' + k2: prepare_mxnet_array(v2[0], v2[1], dtype) for k, v in kernel_params.items() for k2, v2 in v.items()} K_XX_mx = kern.K(mx.nd, X=X_mx, **kernel_params_mx) K_XX2_mx = kern.K(mx.nd, X=X_mx, X2=X2_mx, **kernel_params_mx) Kdiag_mx = kern.Kdiag(mx.nd, X=X_mx, **kernel_params_mx) kern_gpy = gpy_kernel_create() K_XX_gpy, K_XX2_gpy, Kdiag_gpy = [], [], [] for i in range(num_samples): X_i = X[i] if X_isSamples else X X2_i = X2[i] if X2_isSamples else X2 for k, v in kernel_params.items(): kern_1 = getattr(kern_gpy, k) for k2, v2 in v.items(): setattr(kern_1, k2, v2[0][i] if v2[1] else v2[0]) K_XX_gpy.append(np.expand_dims(kern_gpy.K(X_i), axis=0)) K_XX2_gpy.append(np.expand_dims(kern_gpy.K(X_i, X2_i), axis=0)) Kdiag_gpy.append(np.expand_dims(kern_gpy.Kdiag(X_i), axis=0)) K_XX_gpy = np.vstack(K_XX_gpy) K_XX2_gpy = np.vstack(K_XX2_gpy) Kdiag_gpy = np.vstack(Kdiag_gpy) assert np.issubdtype(K_XX_mx.dtype, dtype) assert np.issubdtype(K_XX2_mx.dtype, dtype) assert np.issubdtype(Kdiag_mx.dtype, dtype) assert np.allclose(K_XX_gpy, K_XX_mx.asnumpy()) assert np.allclose(K_XX2_gpy, K_XX2_mx.asnumpy()) assert np.allclose(Kdiag_gpy, Kdiag_mx.asnumpy())
def test_kernel_as_MXFusionFunction(self, dtype, X, X_isSamples, X2, X2_isSamples, lengthscale, lengthscale_isSamples, variance, variance_isSamples, num_samples, input_dim, ARD): X_mx = prepare_mxnet_array(X, X_isSamples, dtype) X2_mx = prepare_mxnet_array(X2, X2_isSamples, dtype) var_mx = prepare_mxnet_array(variance, variance_isSamples, dtype) l_mx = prepare_mxnet_array(lengthscale, lengthscale_isSamples, dtype) X_mf = Variable(shape=X.shape) l_mf = Variable(shape=lengthscale.shape) var_mf = Variable(shape=variance.shape) rbf = RBF(input_dim, ARD, 1., 1., 'rbf', None, dtype) eval = rbf(X_mf, rbf_lengthscale=l_mf, rbf_variance=var_mf).factor variables = {eval.X.uuid: X_mx, eval.rbf_lengthscale.uuid: l_mx, eval.rbf_variance.uuid: var_mx} res_eval = eval.eval(F=mx.nd, variables=variables) kernel_params = rbf.fetch_parameters(variables) res_direct = rbf.K(F=mx.nd, X=X_mx, **kernel_params) assert np.allclose(res_eval.asnumpy(), res_direct.asnumpy()) X_mf = Variable(shape=X.shape) X2_mf = Variable(shape=X2.shape) l_mf = Variable(shape=lengthscale.shape) var_mf = Variable(shape=variance.shape) rbf = RBF(input_dim, ARD, 1., 1., 'rbf', None, dtype) eval = rbf(X_mf, X2_mf, rbf_lengthscale=l_mf, rbf_variance=var_mf).factor variables = {eval.X.uuid: X_mx, eval.X2.uuid: X2_mx, eval.rbf_lengthscale.uuid: l_mx, eval.rbf_variance.uuid: var_mx} res_eval = eval.eval(F=mx.nd, variables=variables) kernel_params = rbf.fetch_parameters(variables) res_direct = rbf.K(F=mx.nd, X=X_mx, X2=X2_mx, **kernel_params) assert np.allclose(res_eval.asnumpy(), res_direct.asnumpy())
def test_clone_gp(self, dtype, X, X_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv, rv_isSamples, num_samples): X_mx = prepare_mxnet_array(X, X_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) rv_mx = prepare_mxnet_array(rv, rv_isSamples, dtype) rv_shape = rv.shape[1:] if rv_isSamples else rv.shape rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) m = Model() m.X_var = Variable(shape=(5, 2)) m.Y = GaussianProcess.define_variable(X=m.X_var, kernel=rbf, shape=rv_shape, dtype=dtype) gp = m.clone().Y.factor variables = { gp.X.uuid: X_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx, gp.random_variable.uuid: rv_mx } log_pdf_rt = gp.log_pdf(F=mx.nd, variables=variables).asnumpy() log_pdf_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rv_i = rv[i] if rv_isSamples else rv rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) log_pdf_np.append( multivariate_normal.logpdf(rv_i[:, 0], mean=None, cov=K_np)) log_pdf_np = np.array(log_pdf_np) isSamples_any = any([ X_isSamples, rbf_lengthscale_isSamples, rbf_variance_isSamples, rv_isSamples ]) 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 assert np.allclose(log_pdf_np, log_pdf_rt)
def test_draw_samples(self, dtype, X, X_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv_shape, num_samples): X_mx = prepare_mxnet_array(X, X_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) rand = np.random.randn(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand.flatten(), dtype=dtype)) rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) X_var = Variable(shape=(5, 2)) gp = GaussianProcess.define_variable(X=X_var, kernel=rbf, shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = { gp.X.uuid: X_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx } samples_rt = gp.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples).asnumpy() samples_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rand_i = rand[i] rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) L_np = np.linalg.cholesky(K_np) sample_np = L_np.dot(rand_i) samples_np.append(sample_np) samples_np = np.array(samples_np) assert np.issubdtype(samples_rt.dtype, dtype) assert get_num_samples(mx.nd, samples_rt) == num_samples assert np.allclose(samples_np, samples_rt)
def test_draw_samples_w_mean(self, dtype, X, X_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv_shape, num_samples): net = nn.HybridSequential(prefix='nn_') with net.name_scope(): net.add( nn.Dense(rv_shape[-1], flatten=False, activation="tanh", in_units=X.shape[-1], dtype=dtype)) net.initialize(mx.init.Xavier(magnitude=3)) X_mx = prepare_mxnet_array(X, X_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) mean_mx = net(X_mx) mean_np = mean_mx.asnumpy() rand = np.random.randn(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand.flatten(), dtype=dtype)) rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) X_var = Variable(shape=(5, 2)) mean_func = MXFusionGluonFunction(net, num_outputs=1, broadcastable=True) mean_var = mean_func(X_var) gp = GaussianProcess.define_variable(X=X_var, kernel=rbf, shape=rv_shape, mean=mean_var, dtype=dtype, rand_gen=rand_gen).factor variables = { gp.X.uuid: X_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx, gp.mean.uuid: mean_mx } samples_rt = gp.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples).asnumpy() samples_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rand_i = rand[i] rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) L_np = np.linalg.cholesky(K_np) sample_np = L_np.dot(rand_i) samples_np.append(sample_np) samples_np = np.array(samples_np) + mean_np assert np.issubdtype(samples_rt.dtype, dtype) assert get_num_samples(mx.nd, samples_rt) == num_samples assert np.allclose(samples_np, samples_rt)
def test_log_pdf_w_mean(self, dtype, X, X_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv, rv_isSamples, num_samples): net = nn.HybridSequential(prefix='nn_') with net.name_scope(): net.add( nn.Dense(rv.shape[-1], flatten=False, activation="tanh", in_units=X.shape[-1], dtype=dtype)) net.initialize(mx.init.Xavier(magnitude=3)) X_mx = prepare_mxnet_array(X, X_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) rv_mx = prepare_mxnet_array(rv, rv_isSamples, dtype) rv_shape = rv.shape[1:] if rv_isSamples else rv.shape mean_mx = net(X_mx) mean_np = mean_mx.asnumpy() rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) X_var = Variable(shape=(5, 2)) mean_func = MXFusionGluonFunction(net, num_outputs=1, broadcastable=True) mean_var = mean_func(X_var) gp = GaussianProcess.define_variable(X=X_var, kernel=rbf, shape=rv_shape, mean=mean_var, dtype=dtype).factor variables = { gp.X.uuid: X_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx, gp.random_variable.uuid: rv_mx, gp.mean.uuid: mean_mx } log_pdf_rt = gp.log_pdf(F=mx.nd, variables=variables).asnumpy() log_pdf_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rv_i = rv[i] if rv_isSamples else rv rv_i = rv_i - mean_np[i] if X_isSamples else rv_i - mean_np[0] rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) log_pdf_np.append( multivariate_normal.logpdf(rv_i[:, 0], mean=None, cov=K_np)) log_pdf_np = np.array(log_pdf_np) isSamples_any = any([ X_isSamples, rbf_lengthscale_isSamples, rbf_variance_isSamples, rv_isSamples ]) 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 assert np.allclose(log_pdf_np, log_pdf_rt)
def test_log_pdf(self, dtype, X, X_isSamples, X_cond, X_cond_isSamples, Y_cond, Y_cond_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv, rv_isSamples, num_samples): from scipy.linalg.lapack import dtrtrs X_mx = prepare_mxnet_array(X, X_isSamples, dtype) X_cond_mx = prepare_mxnet_array(X_cond, X_cond_isSamples, dtype) Y_cond_mx = prepare_mxnet_array(Y_cond, Y_cond_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) rv_mx = prepare_mxnet_array(rv, rv_isSamples, dtype) rv_shape = rv.shape[1:] if rv_isSamples else rv.shape rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) X_var = Variable(shape=(5, 2)) X_cond_var = Variable(shape=(8, 2)) Y_cond_var = Variable(shape=(8, 1)) gp = ConditionalGaussianProcess.define_variable(X=X_var, X_cond=X_cond_var, Y_cond=Y_cond_var, kernel=rbf, shape=rv_shape, dtype=dtype).factor variables = { gp.X.uuid: X_mx, gp.X_cond.uuid: X_cond_mx, gp.Y_cond.uuid: Y_cond_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx, gp.random_variable.uuid: rv_mx } log_pdf_rt = gp.log_pdf(F=mx.nd, variables=variables).asnumpy() log_pdf_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X X_cond_i = X_cond[i] if X_cond_isSamples else X_cond Y_cond_i = Y_cond[i] if Y_cond_isSamples else Y_cond lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rv_i = rv[i] if rv_isSamples else rv rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) Kc_np = rbf_np.K(X_cond_i, X_i) Kcc_np = rbf_np.K(X_cond_i) L = np.linalg.cholesky(Kcc_np) LInvY = dtrtrs(L, Y_cond_i, lower=1, trans=0)[0] LinvKxt = dtrtrs(L, Kc_np, lower=1, trans=0)[0] mu = LinvKxt.T.dot(LInvY) cov = K_np - LinvKxt.T.dot(LinvKxt) log_pdf_np.append( multivariate_normal.logpdf(rv_i[:, 0], mean=mu[:, 0], cov=cov)) log_pdf_np = np.array(log_pdf_np) isSamples_any = any([ X_isSamples, rbf_lengthscale_isSamples, rbf_variance_isSamples, rv_isSamples ]) 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 assert np.allclose(log_pdf_np, log_pdf_rt)
def test_draw_samples_w_mean(self, dtype, X, X_isSamples, X_cond, X_cond_isSamples, Y_cond, Y_cond_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv_shape, num_samples): net = nn.HybridSequential(prefix='nn_') with net.name_scope(): net.add( nn.Dense(rv_shape[-1], flatten=False, activation="tanh", in_units=X.shape[-1], dtype=dtype)) net.initialize(mx.init.Xavier(magnitude=3)) from scipy.linalg.lapack import dtrtrs X_mx = prepare_mxnet_array(X, X_isSamples, dtype) X_cond_mx = prepare_mxnet_array(X_cond, X_cond_isSamples, dtype) Y_cond_mx = prepare_mxnet_array(Y_cond, Y_cond_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) mean_mx = net(X_mx) mean_np = mean_mx.asnumpy() mean_cond_mx = net(X_cond_mx) mean_cond_np = mean_cond_mx.asnumpy() rand = np.random.randn(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand.flatten(), dtype=dtype)) rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) X_var = Variable(shape=(5, 2)) X_cond_var = Variable(shape=(8, 2)) Y_cond_var = Variable(shape=(8, 1)) mean_func = MXFusionGluonFunction(net, num_outputs=1, broadcastable=True) mean_var = mean_func(X_var) mean_cond_var = mean_func(X_cond_var) gp = ConditionalGaussianProcess.define_variable( X=X_var, X_cond=X_cond_var, Y_cond=Y_cond_var, mean=mean_var, mean_cond=mean_cond_var, kernel=rbf, shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = { gp.X.uuid: X_mx, gp.X_cond.uuid: X_cond_mx, gp.Y_cond.uuid: Y_cond_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx, gp.mean.uuid: mean_mx, gp.mean_cond.uuid: mean_cond_mx } samples_rt = gp.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples).asnumpy() samples_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X X_cond_i = X_cond[i] if X_cond_isSamples else X_cond Y_cond_i = Y_cond[i] if Y_cond_isSamples else Y_cond Y_cond_i = Y_cond_i - mean_cond_np[ i] if X_cond_isSamples else Y_cond_i - mean_cond_np[0] lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rand_i = rand[i] rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) Kc_np = rbf_np.K(X_cond_i, X_i) Kcc_np = rbf_np.K(X_cond_i) L = np.linalg.cholesky(Kcc_np) LInvY = dtrtrs(L, Y_cond_i, lower=1, trans=0)[0] LinvKxt = dtrtrs(L, Kc_np, lower=1, trans=0)[0] mu = LinvKxt.T.dot(LInvY) cov = K_np - LinvKxt.T.dot(LinvKxt) L_cov_np = np.linalg.cholesky(cov) sample_np = mu + L_cov_np.dot(rand_i) samples_np.append(sample_np) samples_np = np.array(samples_np) + mean_np assert np.issubdtype(samples_rt.dtype, dtype) assert get_num_samples(mx.nd, samples_rt) == num_samples assert np.allclose(samples_np, samples_rt)
def test_draw_samples(self, dtype, X, X_isSamples, X_cond, X_cond_isSamples, Y_cond, Y_cond_isSamples, rbf_lengthscale, rbf_lengthscale_isSamples, rbf_variance, rbf_variance_isSamples, rv_shape, num_samples): from scipy.linalg.lapack import dtrtrs X_mx = prepare_mxnet_array(X, X_isSamples, dtype) X_cond_mx = prepare_mxnet_array(X_cond, X_cond_isSamples, dtype) Y_cond_mx = prepare_mxnet_array(Y_cond, Y_cond_isSamples, dtype) rbf_lengthscale_mx = prepare_mxnet_array(rbf_lengthscale, rbf_lengthscale_isSamples, dtype) rbf_variance_mx = prepare_mxnet_array(rbf_variance, rbf_variance_isSamples, dtype) rand = np.random.randn(num_samples, *rv_shape) rand_gen = MockMXNetRandomGenerator( mx.nd.array(rand.flatten(), dtype=dtype)) rbf = RBF(2, True, 1., 1., 'rbf', None, dtype) X_var = Variable(shape=(5, 2)) X_cond_var = Variable(shape=(8, 2)) Y_cond_var = Variable(shape=(8, 1)) gp = ConditionalGaussianProcess.define_variable( X=X_var, X_cond=X_cond_var, Y_cond=Y_cond_var, kernel=rbf, shape=rv_shape, dtype=dtype, rand_gen=rand_gen).factor variables = { gp.X.uuid: X_mx, gp.X_cond.uuid: X_cond_mx, gp.Y_cond.uuid: Y_cond_mx, gp.rbf_lengthscale.uuid: rbf_lengthscale_mx, gp.rbf_variance.uuid: rbf_variance_mx } samples_rt = gp.draw_samples(F=mx.nd, variables=variables, num_samples=num_samples).asnumpy() samples_np = [] for i in range(num_samples): X_i = X[i] if X_isSamples else X X_cond_i = X_cond[i] if X_cond_isSamples else X_cond Y_cond_i = Y_cond[i] if Y_cond_isSamples else Y_cond lengthscale_i = rbf_lengthscale[ i] if rbf_lengthscale_isSamples else rbf_lengthscale variance_i = rbf_variance[ i] if rbf_variance_isSamples else rbf_variance rand_i = rand[i] rbf_np = GPy.kern.RBF(input_dim=2, ARD=True) rbf_np.lengthscale = lengthscale_i rbf_np.variance = variance_i K_np = rbf_np.K(X_i) Kc_np = rbf_np.K(X_cond_i, X_i) Kcc_np = rbf_np.K(X_cond_i) L = np.linalg.cholesky(Kcc_np) LInvY = dtrtrs(L, Y_cond_i, lower=1, trans=0)[0] LinvKxt = dtrtrs(L, Kc_np, lower=1, trans=0)[0] mu = LinvKxt.T.dot(LInvY) cov = K_np - LinvKxt.T.dot(LinvKxt) L_cov_np = np.linalg.cholesky(cov) sample_np = mu + L_cov_np.dot(rand_i) samples_np.append(sample_np) samples_np = np.array(samples_np) assert np.issubdtype(samples_rt.dtype, dtype) assert get_num_samples(mx.nd, samples_rt) == num_samples assert np.allclose(samples_np, samples_rt)