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

… unravel (apache#15048)

python/mxnet/test_utils.py

@@ -419,6 +419,12 @@ def rand_shape_nd(num_dim, dim=10):
     return tuple(rnd.randint(1, dim+1, size=num_dim))
 
 
+def rand_coord_2d(x_low, x_high, y_low, y_high):
+    x = np.random.randint(x_low, x_high, dtype=np.int64)
+    y = np.random.randint(y_low, y_high, dtype=np.int64)
+    return x, y
+
+
 def np_reduce(dat, axis, keepdims, numpy_reduce_func):
     """Compatible reduce for old version of NumPy.
 

tests/nightly/test_large_array.py

@@ -15,9 +15,9 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import mxnet as mx
 import numpy as np
-from mxnet.test_utils import rand_ndarray, assert_almost_equal
+import mxnet as mx
+from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d
 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)
@@ -215,9 +208,9 @@ 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)
+    a = mx.nd.ones(shape=(256 * 35, 1024 * 1024))
+    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,33 @@ 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))
 
 
+@with_seed()
+def test_ravel_multi_index():
+    x1, y1 = rand_coord_2d((LARGE_X - 100), LARGE_X, 10, SMALL_Y)
+    x2, y2 = rand_coord_2d((LARGE_X - 200), LARGE_X, 9, SMALL_Y)
+    x3, y3 = rand_coord_2d((LARGE_X - 300), LARGE_X, 8, SMALL_Y)
+    indices_2d = [[x1, x2, x3], [y1, y2, y3]]
+    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
+
+
+@with_seed()
+def test_unravel_index():
+    x1, y1 = rand_coord_2d((LARGE_X - 100), LARGE_X, 10, SMALL_Y)
+    x2, y2 = rand_coord_2d((LARGE_X - 200), LARGE_X, 9, SMALL_Y)
+    x3, y3 = rand_coord_2d((LARGE_X - 300), LARGE_X, 8, SMALL_Y)
+    original_2d_indices = [[x1, x2, x3], [y1, y2, y3]]
+    idx_numpy = np.ravel_multi_index(original_2d_indices, (LARGE_X, SMALL_Y))
+    indices_2d = mx.nd.unravel_index(mx.nd.array(idx_numpy, 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()