def test_topk(): b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) k = nd.topk(b, k=10, axis=0, dtype=np.int64) assert np.sum(k.asnumpy() == (LARGE_X - 1)) == SMALL_Y ind, val = mx.nd.topk(b, k=3, axis=0, dtype=np.int64, ret_typ="both", is_ascend=False) assert np.all(ind == val) b = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) l = nd.topk(b, k=1, axis=-1, dtype=np.int64, ret_typ="value") assert l.sum() == np.sum(np.arange(0, SMALL_Y))
def test_repeat(): x = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X//2) y = nd.repeat(x, repeats=2, axis = 1) assert y.shape == (SMALL_Y, LARGE_X) assert y[0][1] == 0 assert y[-1][-1] == SMALL_Y-1 x = create_2d_tensor(rows=SMALL_Y//2, columns=LARGE_X) y = nd.repeat(x, repeats=2, axis = 0) assert y.shape == (SMALL_Y, LARGE_X) assert y[0][1] == 0 assert y[-1][0] == SMALL_Y//2-1
def test_copy_to(): a = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) b = nd.array(np.zeros((SMALL_Y, LARGE_X))) c = a.copyto(b) assert c is b print(b) assert b[0][-1] == LARGE_X-1
def test_flatten(): a = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y).reshape( (LARGE_X // 2, 2, SMALL_Y)) b = nd.flatten(a) assert b[-1][-1] == (LARGE_X - 1) assert b[-1][0] == (LARGE_X - 2) assert b.shape == (LARGE_X // 2, SMALL_Y * 2)
def test_pad(): x = create_2d_tensor(rows=SMALL_Y-2, columns=LARGE_X//2-2, dtype=np.float32).reshape(1 , 1, SMALL_Y-2, LARGE_X//2-2) y = nd.pad(x, mode="edge", pad_width=(0, 0, 0, 0, 1, 1, 1, 1)) assert y[0][0][1][0] == 0 assert y[0][0][1][-1] == 0 assert y[0][0][-1][0] == SMALL_Y-3 assert y[0][0][-1][-1] == SMALL_Y-3 assert y.shape == (1, 1, SMALL_Y, LARGE_X//2)
def test_slice_axis(): a = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) c = nd.slice_axis(a, axis=0, begin=0, end=SMALL_Y//2) d = nd.slice_axis(a, axis=1, begin=0, end=LARGE_X//2) assert c.shape[0] == a.shape[0]//2 assert d.shape[1] == a.shape[1]//2 assert c[-1][0] == (SMALL_Y//2-1) assert d[-1][-1] == (SMALL_Y-1)
def test_load_save(): x = create_2d_tensor(SMALL_Y, LARGE_X) tmp = tempfile.mkdtemp() tmpfile = os.path.join(tmp, 'large_tensor') nd.save(tmpfile, [x]) y = nd.load(tmpfile) y = y[0] assert x[0][0] == y[0][0] assert x[-1][-1]== y[-1][-1]
def test_slice_like(): a = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) b = nd.array(np.ones((SMALL_Y//2, LARGE_X//2))) c = nd.slice_like(a, b) d = nd.slice_like(a, b, axes=(0)) e = nd.slice_like(a, b, axes=(-1)) assert c.shape == b.shape assert d.shape[0] == b.shape[0] assert e.shape[-1] == b.shape[-1] assert c[0][-1] == 0 assert d[-1][0] == (SMALL_Y//2-1) assert e[-1][-1] == (SMALL_Y-1)
def test_size(): b = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) mx.nd.waitall() assert b.size == LARGE_SIZE
def test_shape(): b = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) mx.nd.waitall() assert b.shape == (SMALL_Y, LARGE_X)
def test_cast(): x = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) y = nd.cast(x, np.int32) assert y.dtype == np.int32 assert y[-1][-1] == SMALL_Y-1
def test_sort(): b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) s = nd.sort(b, axis=0, is_ascend=False) assert np.sum(s[-1][SMALL_Y//2:SMALL_Y].asnumpy() == 0).all() s = nd.sort(b, is_ascend=False) assert np.sum(s[0].asnumpy() == 0).all()
def test_argsort(): b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) s = nd.argsort(b, axis=0, is_ascend=False, dtype=np.int64) mx.nd.waitall() assert (s[0].asnumpy() == (LARGE_X - 1)).all()
def test_flip(): b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) t = nd.flip(b, axis=0) assert np.sum(t[-1, :].asnumpy() == 0) == b.shape[1] assert t.shape == (LARGE_X, SMALL_Y)
def test_swapaxes(): b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) t = nd.swapaxes(b, dim1=0, dim2=1) assert np.sum(t[:, -1].asnumpy() == (LARGE_X - 1)) == b.shape[1] assert t.shape == (SMALL_Y, LARGE_X)
def test_astype(): x = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) y = x.astype('int32') assert y.dtype == np.int32 assert y[-1][-1] == SMALL_Y-1
def test_broadcast_axes(): a = create_2d_tensor(rows=1, columns=LARGE_X) b = nd.broadcast_axis(a, axis=[0], size=2) assert b.shape == (a.shape[0]*2, a.shape[1])
def test_max(): a = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) b = nd.max(a, axis=0) assert b[0] == (SMALL_Y-1) assert b[-1] == (SMALL_Y-1)
def test_min(): a = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) b = nd.min(a, axis=0) assert b[0] == 0 assert b[-1] == 0
def test_mean(): a = create_2d_tensor(rows=SMALL_Y, columns=LARGE_X) b = nd.mean(a, axis=0) assert b[0] == (SMALL_Y/2-1)
def test_transpose(): b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y) t = b.T assert np.sum(t[:, -1].asnumpy() == (LARGE_X - 1)) == b.shape[1] assert t.shape == (SMALL_Y, LARGE_X)