From 3a430daaba4ea04bae90e2326e66febfef70a00f Mon Sep 17 00:00:00 2001 From: Nick Guletskii Date: Tue, 14 May 2019 00:36:41 +0300 Subject: [PATCH] Move the index_array tests into a single function (test_index_array) --- tests/python/unittest/test_operator.py | 124 +++++++++++++------------ 1 file changed, 64 insertions(+), 60 deletions(-) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 1479c4facbca..f36d46af9f01 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -8055,67 +8055,71 @@ def test_image_normalize(): check_numeric_gradient(img_norm_sym, [data_in_4d], atol=0.001) @with_seed() -def test_index_array_default(): - for shape in [(10,), (7, 5, 29), (5, 7, 11, 13, 17, 19)]: - data = mx.symbol.Variable("data") - index_array = mx.sym.contrib.index_array(data) +def test_index_array(): + def test_index_array_default(): + for shape in [(10,), (7, 5, 29), (5, 7, 11, 13, 17, 19)]: + data = mx.symbol.Variable("data") + index_array = mx.sym.contrib.index_array(data) + + input_array = np.ones(shape) + mgrid = np.mgrid[tuple(slice(0, x) for x in shape)] + expected = np.stack(mgrid, axis=-1) + + check_symbolic_forward(index_array, [input_array], [expected]) + check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) + + def test_index_array_default_zero_dim(): + with mx.np_compat(active=True): + data = mx.symbol.Variable("data") + index_array = mx.sym.contrib.index_array(data) + + input_array = np.ones(()) + expected = np.zeros((0,)) + + check_symbolic_forward(index_array, [input_array], [expected]) + check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) + + def test_index_array_default_zero_size(): + with mx.np_compat(active=True): + data = mx.symbol.Variable("data") + index_array = mx.sym.contrib.index_array(data) + + input_array = np.ones((0, 0, 0)) + expected = np.zeros((0, 0, 0, 3)) + + check_symbolic_forward(index_array, [input_array], [expected]) + check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) + + def test_index_array_select_axes(): + shape = (5, 7, 11, 13, 17, 19) + for axes in [(3,), (4, 1), (5, 1, 3), (-1,), (-5, -1, -3)]: + data = mx.symbol.Variable("data") + index_array = mx.sym.contrib.index_array(data, axes=axes) + + input_array = np.ones(shape) + mgrid = np.mgrid[tuple(slice(0, x) for x in shape)] + expected = np.stack(mgrid, axis=-1)[..., axes] + + check_symbolic_forward(index_array, [input_array], [expected]) + check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) + + def test_index_array_select_axes_zero_size(): + with mx.np_compat(active=True): + data = mx.symbol.Variable("data") + index_array = mx.sym.contrib.index_array(data, axes=(2, 1)) + + input_array = np.ones((0, 0, 0, 0)) + expected = np.zeros((0, 0, 2)) + + check_symbolic_forward(index_array, [input_array], [expected]) + check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) + + test_index_array_default() + test_index_array_default_zero_dim() + test_index_array_default_zero_size() + test_index_array_select_axes() + test_index_array_select_axes_zero_size() - input_array = np.ones(shape) - mgrid = np.mgrid[tuple(slice(0, x) for x in shape)] - expected = np.stack(mgrid, axis=-1) - - check_symbolic_forward(index_array, [input_array], [expected]) - check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) - -@with_seed() -def test_index_array_default_zero_dim(): - with mx.np_compat(active=True): - data = mx.symbol.Variable("data") - index_array = mx.sym.contrib.index_array(data) - - input_array = np.ones(()) - expected = np.zeros((0,)) - - check_symbolic_forward(index_array, [input_array], [expected]) - check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) - -@with_seed() -def test_index_array_default_zero_size(): - with mx.np_compat(active=True): - data = mx.symbol.Variable("data") - index_array = mx.sym.contrib.index_array(data) - - input_array = np.ones((0, 0, 0)) - expected = np.zeros((0, 0, 0, 3)) - - check_symbolic_forward(index_array, [input_array], [expected]) - check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) - -@with_seed() -def test_index_array_select_axes(): - shape = (5, 7, 11, 13, 17, 19) - for axes in [(3,), (4, 1), (5, 1, 3), (-1,), (-5, -1, -3)]: - data = mx.symbol.Variable("data") - index_array = mx.sym.contrib.index_array(data, axes=axes) - - input_array = np.ones(shape) - mgrid = np.mgrid[tuple(slice(0, x) for x in shape)] - expected = np.stack(mgrid, axis=-1)[..., axes] - - check_symbolic_forward(index_array, [input_array], [expected]) - check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) - -@with_seed() -def test_index_array_select_axes_zero_size(): - with mx.np_compat(active=True): - data = mx.symbol.Variable("data") - index_array = mx.sym.contrib.index_array(data, axes=(2, 1)) - - input_array = np.ones((0, 0, 0, 0)) - expected = np.zeros((0, 0, 2)) - - check_symbolic_forward(index_array, [input_array], [expected]) - check_symbolic_backward(index_array, [input_array], [np.ones(expected.shape)], [np.zeros_like(input_array)]) @with_seed() def test_scalar_tensor_creation():