[MXNET-1401] adding more operators to test support for Large Tensor

tests/nightly/test_large_array.py

@@ -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))
@@ -128,6 +138,38 @@ def test_take():
  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))
@@ -175,7 +217,6 @@ def test_where():
  b = nd.broadcast_to(b, shape=(b.shape[0], 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]
@@ -186,7 +227,8 @@ def test_pick():
  b = mx.nd.ones(shape=(256*35,))
  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]
@@ -202,6 +244,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]
@@ -217,6 +260,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).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()