diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py index 37d867a2d576..9773adeac87b 100644 --- a/tests/python/gpu/test_operator_gpu.py +++ b/tests/python/gpu/test_operator_gpu.py @@ -2143,6 +2143,30 @@ def test_index_array_default(): 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) @@ -2157,6 +2181,18 @@ def test_index_array_select_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)]) + # isolated execution bulking test function to be invoked with different env var settings def _test_bulking_in_process(seed, time_per_iteration): data_shape = (10,) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 42f3a9a629a3..1479c4facbca 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -8067,6 +8067,30 @@ def test_index_array_default(): 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) @@ -8081,6 +8105,18 @@ def test_index_array_select_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(): assertRaises(MXNetError, mx.nd.zeros, shape=())