def test_slice_like():
a = create_vector(size=LARGE_X)
b = nd.ones(LARGE_X // 2)
c = nd.slice_like(a, b)
assert c.shape == b.shape
assert c[0] == 0
assert c[-1] == (LARGE_X // 2 - 1)
def test_copy_to():
a = create_vector(size=LARGE_X)
# keeping dtype same as input uses parallel copy which is much faster
b = nd.zeros(LARGE_X, dtype=np.int64)
c = a.copyto(b)
assert c is b
assert b[-1] == LARGE_X - 1
assert b[0] == 0
def check_load_save():
x = create_vector(size=LARGE_X)
with TemporaryDirectory() as tmp:
tmpfile = os.path.join(tmp, 'large_vector')
nd.save(tmpfile, [x])
y = nd.load(tmpfile)
y = y[0]
assert x[0] == y[0]
assert x[-1] == y[-1]
def check_topk():
a = create_vector(size=LARGE_X)
ind = nd.topk(a, k=10, axis=0, dtype=np.int64)
for i in range(10):
assert ind[i] == (LARGE_X - i - 1)
ind, val = mx.nd.topk(a, k=3, axis=0, dtype=np.int64, ret_typ="both", is_ascend=False)
assert np.all(ind == val)
val = nd.topk(a, k=1, axis=0, dtype=np.int64, ret_typ="value")
assert val == (LARGE_X - 1)
def test_load_save():
x = create_vector(size=LARGE_X)
tmp = tempfile.mkdtemp()
tmpfile = os.path.join(tmp, 'large_vector')
nd.save(tmpfile, [x])
y = nd.load(tmpfile)
y = y[0]
assert x[0] == y[0]
assert x[-1] == y[-1]
def test_sort():
a = create_vector(size=LARGE_X)
def test_descend(x):
s = nd.sort(x, axis=0, is_ascend=False)
assert s[-1] == 0
def test_ascend(x):
s = nd.sort(x, is_ascend=True)
assert s[0] == 0
test_descend(a)
test_ascend(a)
def test_topk():
b = create_vector(size=LARGE_X)
ind = nd.topk(b, k=10, axis=0, dtype=np.int64)
assert np.sum(ind.asnumpy() == (LARGE_X - 1)) == 1
ind, val = mx.nd.topk(b,
k=3,
axis=0,
dtype=np.int64,
ret_typ="both",
is_ascend=False)
assert np.all(ind == val)
val = nd.topk(b, k=1, axis=0, dtype=np.int64, ret_typ="value")
assert val.sum() == (LARGE_X - 1)
def test_size():
b = create_vector(size=LARGE_X)
#explicit wait_to_read()
assert b[0] == 0
assert b.size == LARGE_X
def test_argsort():
a = create_vector(size=LARGE_X)
s = nd.argsort(a, axis=0, is_ascend=False, dtype=np.int64)
assert s[0] == (LARGE_X - 1)
def test_sort():
b = create_vector(size=LARGE_X)
s = nd.sort(b, axis=0, is_ascend=False)
assert np.sum(s[-1].asnumpy() == 0).all()
s = nd.sort(b, is_ascend=True)
assert np.sum(s[0].asnumpy() == 0).all()
def test_clip():
a = create_vector(LARGE_X)
res = nd.clip(a, a_min=100, a_max=1000)
assert res[-1] == 1000
def test_astype():
x = create_vector(size=LARGE_X // 4)
x = nd.tile(x, 4)
y = x.astype('int32')
assert y.dtype == np.int32
assert y[-1] == LARGE_X // 4 - 1
def check_shape():
b = create_vector(size=LARGE_X)
# explicit wait_to_read()
assert b[0] == 0
assert b.shape[0] == LARGE_X
def test_clip():
a = create_vector(LARGE_X)
res = nd.clip(a, a_min=100, a_max=1000)
assert np.sum(res[-1].asnumpy() == 1000) == 1
def test_min():
a = create_vector(size=LARGE_X)
b = nd.min(a, axis=0)
assert b[0] == 0
assert b[-1] == 0
def test_max():
a = create_vector(size=LARGE_X)
b = nd.max(a, axis=0)
assert b[0] == (LARGE_X - 1)
def test_cast():
x = create_vector(size=LARGE_X // 4)
x = nd.tile(x, 4)
y = nd.cast(x, np.int32)
assert y.dtype == np.int32
assert y[-1] == LARGE_X // 4 - 1
def test_slice_axis():
a = create_vector(size=LARGE_X)
med = LARGE_X // 2
c = nd.slice_axis(a, axis=0, begin=0, end=med)
assert c.shape[0] == a.shape[0] // 2
assert c[-1][0] == (med - 1)
def test_repeat():
x = create_vector(size=LARGE_X // 2)
y = nd.repeat(x, repeats=2, axis=0)
assert y.shape[0] == LARGE_X
assert y[1] == 0
assert y[LARGE_X - 1] == LARGE_X // 2 - 1
def test_argmin():
a = create_vector(LARGE_X, dtype=np.float32)
assert a[0] == 0
idx = mx.nd.argmin(a, axis=0)
assert idx[0] == 0
assert idx.shape[0] == 1
def test_argsort():
b = create_vector(size=LARGE_X)
s = nd.argsort(b, axis=0, is_ascend=False, dtype=np.int64)
assert (s[0].asnumpy() == (LARGE_X - 1)).all()
