[MXNET-1408] Adding test to verify Large Tensor Support for ravel and unravel

Author: Rohit Kumar Srivastava

tests/nightly/test_large_array.py

@@ -15,8 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mxnet as mx
 import numpy as np
+import mxnet as mx
 from mxnet.test_utils import rand_ndarray, assert_almost_equal
 from mxnet import gluon, nd
 from tests.python.unittest.common import with_seed
@@ -74,7 +74,7 @@ def test_ndarray_random_randint():
     # check if randint can generate value greater than 2**32 (large)
     low_large_value = 2**32
     high_large_value = 2**34
-    a = nd.random.randint(low_large_value,high_large_value, dtype=np.int64)
+    a = nd.random.randint(low_large_value, high_large_value, dtype=np.int64)
     low = mx.nd.array([low_large_value], dtype='int64')
     high = mx.nd.array([high_large_value], dtype='int64')
     assert a.__gt__(low) and a.__lt__(high)
@@ -130,13 +130,6 @@ def test_clip():
     assert np.sum(res[-1].asnumpy() == 1000) == a.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 * SMALL_Y).reshape(LARGE_X, SMALL_Y)
     outs = nd.split(a, num_outputs=SMALL_Y, axis=1)
@@ -216,8 +209,8 @@ def test_where():
 
 def test_pick():
     a = mx.nd.ones(shape=(256*35, 1024*1024))
-    b = mx.nd.ones(shape=(256*35,))
-    res = mx.nd.pick(a,b)
+    b = mx.nd.ones(shape=(256*35, ))
+    res = mx.nd.pick(a, b)
     assert res.shape == b.shape
 
 
@@ -252,11 +245,9 @@ 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)
 
-
+@with_seed()
 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_np = np.random.random((LARGE_X, SMALL_Y)).astype(np.float32)
     a = mx.nd.array(a_np)
 
     # k == 0
@@ -274,11 +265,26 @@ def test_diag():
     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))
+    k = np.random.randint(-min(LARGE_X, SMALL_Y) + 1, min(LARGE_X, SMALL_Y))
     r = mx.nd.diag(a, k=k)
     assert_almost_equal(r.asnumpy(), np.diag(a_np, k=k))
 
 
+def test_ravel_multi_index():
+    indices_2d = [[LARGE_X-1, LARGE_X-100, 6], [SMALL_Y-1, SMALL_Y-10, 1]]
+    idx = mx.nd.ravel_multi_index(mx.nd.array(indices_2d, dtype=np.int64), shape=(LARGE_X, SMALL_Y))
+    idx_numpy = np.ravel_multi_index(indices_2d, (LARGE_X, SMALL_Y))
+    assert np.sum(1 for i in range(idx.size) if idx[i] == idx_numpy[i]) == 3
+
+
+def test_unravel_index():
+    original_2d_indices = [[LARGE_X-1, LARGE_X-100, 6], [SMALL_Y-1, SMALL_Y-10, 1]]
+    idx = mx.nd.ravel_multi_index(mx.nd.array(original_2d_indices, dtype=np.int64), shape=(LARGE_X, SMALL_Y))
+    idx_numpy = np.ravel_multi_index(original_2d_indices, (LARGE_X, SMALL_Y))
+    indices_2d = mx.nd.unravel_index(mx.nd.array(idx, dtype=np.int64), shape=(LARGE_X, SMALL_Y))
+    assert (indices_2d.asnumpy() == np.array(original_2d_indices)).all()
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()