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)
