diff --git a/python/mxnet/test_utils.py b/python/mxnet/test_utils.py index fb40474bc678..6668e5e3f33a 100644 --- a/python/mxnet/test_utils.py +++ b/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_2d_coord(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. diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index 4f47acf9ddd1..66d4227b06b0 100644 --- a/tests/nightly/test_large_array.py +++ b/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_2d_coord 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,34 @@ 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_2d_coord((LARGE_X - 100), LARGE_X, 10, SMALL_Y) + x2, y2 = rand_2d_coord((LARGE_X - 200), LARGE_X, 9, SMALL_Y) + x3, y3 = rand_2d_coord((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_2d_coord((LARGE_X - 100), LARGE_X, 10, SMALL_Y) + x2, y2 = rand_2d_coord((LARGE_X - 200), LARGE_X, 9, SMALL_Y) + x3, y3 = rand_2d_coord((LARGE_X - 300), LARGE_X, 8, SMALL_Y) + original_2d_indices = [[x1, x2, x3], [y1, y2, y3]] + #idx = mx.np.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_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()