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

Commit

Permalink
adding more operators to test support for Large Tensor
Browse files Browse the repository at this point in the history
  • Loading branch information
Rohit Kumar Srivastava committed May 14, 2019
1 parent 3122548 commit 6cb821b
Showing 1 changed file with 62 additions and 20 deletions.
82 changes: 62 additions & 20 deletions tests/nightly/test_large_array.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

import mxnet as mx
import numpy as np
from mxnet.test_utils import rand_ndarray, assert_almost_equal
from mxnet.test_utils import rand_ndarray, assert_almost_equal, check_numeric_gradient
from mxnet import gluon, nd
from tests.python.unittest.common import with_seed

Expand All @@ -28,7 +28,6 @@
SMALL_Y = 50
LARGE_SIZE = LARGE_X * SMALL_Y


def test_gluon_embedding():
m = gluon.nn.Embedding(SMALL_Y, MEDIUM_X)
m.initialize()
Expand All @@ -37,26 +36,31 @@ def test_gluon_embedding():
assert b.shape == (MEDIUM_X, SMALL_Y, MEDIUM_X)
assert b.asnumpy().size == LARGE_SIZE


def test_ndarray_zeros():
a = nd.zeros(shape=(LARGE_X, SMALL_Y))
assert a[-1][0] == 0
assert a.shape == (LARGE_X, SMALL_Y)
assert a.size == LARGE_SIZE


def test_ndarray_ones():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
assert a[-1][0] == 1
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))
assert a[-1][0] != 0


@with_seed()
def test_ndarray_random_randint():
a = nd.random.randint(100, 10000, shape=(LARGE_X, SMALL_Y))
Expand All @@ -69,12 +73,10 @@ def test_ndarray_random_randint():
high = mx.nd.array([high_large_value], dtype='int64')
assert a.__gt__(low) & a.__lt__(high)


def test_ndarray_empty():
a = nd.empty((LARGE_X, SMALL_Y))
assert a.shape == (LARGE_X, SMALL_Y)


def test_elementwise():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.ones(shape=(LARGE_X, SMALL_Y))
Expand All @@ -85,26 +87,22 @@ def test_elementwise():
res = nd.sqrt(a + 3)
assert np.sum(res[-1].asnumpy() == 2) == a.shape[1]


def test_reduce():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
assert nd.sum(a).asnumpy() == a.shape[0] * a.shape[1]


def test_dot():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.ones(shape=(SMALL_Y, SMALL_Y))
res = nd.dot(a, b)
assert np.sum(res[-1].asnumpy() == SMALL_Y) == b.shape[1]


def test_FullyConnected():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.ones(shape=(SMALL_Y, SMALL_Y))
res = nd.FullyConnected(a, b, num_hidden=b.shape[1], no_bias=True)
assert np.sum(res[-1].asnumpy() == SMALL_Y) == b.shape[1]


def test_broadcast():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
Expand All @@ -113,53 +111,74 @@ def test_broadcast():
res = mx.nd.broadcast_like(b, a)
assert np.sum(res[-1].asnumpy() == LARGE_X) == a.shape[1]


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]


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_split():
a = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
b = nd.broadcast_to(a, shape=(a.shape[0], SMALL_Y))
outs = nd.split(b, num_outputs=10, axis=0)
for i, out in enumerate(outs):
assert np.sum(out[0].asnumpy() == i * out.shape[0]) == b.shape[1]
cat = nd.concat(*outs, dim=0)
assert np.sum(cat[-1].asnumpy() == LARGE_X) == cat.shape[1]
stack = nd.stack(*outs)
assert np.sum(stack[-1,-1,:].asnumpy() == LARGE_X) == b.shape[1]

def test_argmin():
a = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
b = nd.broadcast_to(a, shape=(a.shape[0], SMALL_Y))
idx = mx.nd.argmax(b)
assert idx.asnumpy() == (LARGE_X * 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))
assert np.sum(res[-1].asnumpy() == 1) == res.shape[1]


def test_slice_assign():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
a[LARGE_X-1:LARGE_X] = 1000
assert np.sum(a[-1].asnumpy() == 1000) == a.shape[1]


def test_expand_dims():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
res = nd.expand_dims(a, axis=1)
assert res.shape == (a.shape[0], 1, a.shape[1])


def test_squeeze():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
data = nd.expand_dims(a, axis=1)
res = nd.squeeze(data)
assert res.shape == a.shape


def test_broadcast_div():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.ones(shape=(LARGE_X, 1)) * 2
res = a / b
assert np.sum(res[-1].asnumpy() == 0.5) == a.shape[1]


def test_Dense(ctx=mx.cpu(0)):
data = mx.nd.ones(shape=(50*1000*1000, 100))
linear = gluon.nn.Dense(100)
Expand All @@ -168,7 +187,6 @@ def test_Dense(ctx=mx.cpu(0)):
res.wait_to_read()
assert res.shape == (50000000, 100)


def test_where():
a = nd.ones(shape=(LARGE_X, SMALL_Y))
b = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
Expand All @@ -180,7 +198,6 @@ def test_where():
res = nd.sparse.where(csr_cond, a, b)
assert np.sum(res[0].asnumpy() == 1) == b.shape[1]


def test_pick():
a = mx.nd.ones(shape=(256*35, 1024*1024))
b = mx.nd.ones(shape=(256*35,))
Expand Down Expand Up @@ -217,6 +234,31 @@ 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).astype('float32')

# 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 6cb821b

Please sign in to comment.