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

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()