Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
[MXNET-1401] adding more operators to test support for Large Tensor (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
access2rohit authored and apeforest committed May 22, 2019
1 parent 90b66b8 commit e5316b1
Showing 1 changed file with 70 additions and 8 deletions.
78 changes: 70 additions & 8 deletions tests/nightly/test_large_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,16 @@ def test_ndarray_ones():
assert nd.sum(a).asnumpy() == LARGE_SIZE


def test_ndarray_convert():
a = nd.zeros(shape=(LARGE_X, SMALL_Y))
b = a.astype(np.int32)
b.wait_to_read()
assert b.dtype == np.int32
b = a.tostype('row_sparse')
b.wait_to_read()
assert isinstance(b, mx.nd.sparse.RowSparseNDArray)


@with_seed()
def test_ndarray_random_uniform():
a = nd.random.uniform(shape=(LARGE_X, SMALL_Y))
Expand Down Expand Up @@ -115,10 +125,9 @@ def test_broadcast():


def test_clip():
a = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
b = nd.broadcast_to(a, shape=(a.shape[0], SMALL_Y))
res = nd.clip(b, a_min=100, a_max=1000)
assert np.sum(res[-1].asnumpy() == 1000) == b.shape[1]
a = nd.arange(0, LARGE_X * SMALL_Y).reshape(LARGE_X, SMALL_Y)
res = nd.clip(a, a_min=100, a_max=1000)
assert np.sum(res[-1].asnumpy() == 1000) == a.shape[1]


def test_take():
Expand All @@ -128,6 +137,32 @@ def test_take():
assert np.sum(res[-1].asnumpy() == 1) == res.shape[1]


def test_split():
a = nd.arange(0, LARGE_X * SMALL_Y).reshape(LARGE_X, SMALL_Y)
outs = nd.split(a, num_outputs=SMALL_Y, axis=1)
result = sum(1 for i, v in enumerate(outs) if i == v[0].asnumpy())
assert result == a.shape[1]


def test_argmin():
a = nd.arange(0, LARGE_X * SMALL_Y).reshape(LARGE_X, SMALL_Y)
idx = mx.nd.argmin(a, axis=0)
assert idx.shape[0] == SMALL_Y


def test_tile():
a = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
b = nd.tile(a, reps=(1, SMALL_Y))
assert np.sum(b[-1].asnumpy() == LARGE_X) == b.shape[1]


def test_take():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
idx = nd.arange(LARGE_X - 1000, LARGE_X)
res = nd.take(a, idx)
assert np.sum(res[-1].asnumpy() == 1) == res.shape[1]


def test_slice():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
res = nd.slice(a, begin=(LARGE_X-1000, 1), end=(LARGE_X, SMALL_Y))
Expand Down Expand Up @@ -171,14 +206,12 @@ def test_Dense(ctx=mx.cpu(0)):

def test_where():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
b = nd.broadcast_to(b, shape=(b.shape[0], SMALL_Y))
b = nd.arange(0, LARGE_X * SMALL_Y).reshape(LARGE_X, SMALL_Y)
res = nd.where(b > 100, a, b)
assert np.sum(res[-1].asnumpy() == 1) == b.shape[1]

csr_cond = nd.sparse.cast_storage(b < 10, 'csr')
res = nd.sparse.where(csr_cond, a, b)
assert np.sum(res[0].asnumpy() == 1) == b.shape[1]
assert np.sum(res[0].asnumpy() == 1) == 10


def test_pick():
Expand All @@ -187,6 +220,7 @@ def test_pick():
res = mx.nd.pick(a,b)
assert res.shape == b.shape


def test_depthtospace():
def numpy_depth_to_space(x, blocksize):
b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
Expand All @@ -202,6 +236,7 @@ def numpy_depth_to_space(x, blocksize):
output = mx.nd.depth_to_space(data, 2)
assert_almost_equal(output.asnumpy(), expected, atol=1e-3, rtol=1e-3)


def test_spacetodepth():
def numpy_space_to_depth(x, blocksize):
b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
Expand All @@ -217,6 +252,33 @@ def numpy_space_to_depth(x, blocksize):
output = mx.nd.space_to_depth(data, 2)
assert_almost_equal(output.asnumpy(), expected, atol=1e-3, rtol=1e-3)


def test_diag():
h = np.random.randint(2,9)
w = np.random.randint(2,9)
a_np = np.random.random((LARGE_X, 64)).astype(np.float32)
a = mx.nd.array(a_np)

# k == 0
r = mx.nd.diag(a)
assert_almost_equal(r.asnumpy(), np.diag(a_np))

# k == 1
k = 1
r = mx.nd.diag(a, k=k)
assert_almost_equal(r.asnumpy(), np.diag(a_np, k=k))

# k == -1
k = -1
r = mx.nd.diag(a, k=k)
assert_almost_equal(r.asnumpy(), np.diag(a_np, k=k))

# random k
k = np.random.randint(-min(LARGE_X, 64) + 1, min(h, w))
r = mx.nd.diag(a, k=k)
assert_almost_equal(r.asnumpy(), np.diag(a_np, k=k))


if __name__ == '__main__':
import nose
nose.runmodule()

0 comments on commit e5316b1

Please sign in to comment.