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()