From 0b4fd0fa5b03c30f6599dab754be36ef1e04a93a Mon Sep 17 00:00:00 2001 From: JYChen Date: Wed, 15 Nov 2023 11:54:17 +0800 Subject: [PATCH] Support py-boolean in indexing (#58856) * fix single Py_Bool indexing and add unittest case * fix ci error * complex32 -> complex128 --- .../base/dygraph/tensor_patch_methods.py | 4 +- python/paddle/base/variable_index.py | 20 ++-- test/indexing/test_getitem.py | 99 ++++++++++++++++--- test/indexing/test_setitem.py | 90 +++++++++++++++-- 4 files changed, 178 insertions(+), 35 deletions(-) diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py index 05adddbee8bde..a319c03be5843 100644 --- a/python/paddle/base/dygraph/tensor_patch_methods.py +++ b/python/paddle/base/dygraph/tensor_patch_methods.py @@ -870,7 +870,9 @@ def contain_tensor_or_list(item): item = (item,) for slice_item in item: - if isinstance(slice_item, (list, np.ndarray, Variable, range)): + if isinstance( + slice_item, (list, np.ndarray, Variable, range, bool) + ): return True return False diff --git a/python/paddle/base/variable_index.py b/python/paddle/base/variable_index.py index c5264bf482bb1..41aa40c2574e8 100644 --- a/python/paddle/base/variable_index.py +++ b/python/paddle/base/variable_index.py @@ -271,7 +271,7 @@ def replace_none(item): def is_integer_or_scalar_tensor(ele): from .framework import Variable - if isinstance(ele, int): + if type(ele) is int: return True elif isinstance(ele, Variable): # NOTE(zoooo0820): For compatibility, if FLAGS_set_to_1d is set to True, @@ -693,7 +693,8 @@ def parse_index(x, indices): ) estimated_dim = 0 - for dim, slice_item in enumerate(indices): + dim = 0 + for i, slice_item in enumerate(indices): start, end, step = None, None, None if is_integer_or_scalar_tensor(slice_item): if ( @@ -718,21 +719,22 @@ def parse_index(x, indices): start = slice_item step = 1 end = slice_item + 1 if slice_item != -1 else MAX_INTEGER + dim += 1 elif isinstance(slice_item, bool): # single bool is advanced-indexing none_axes.append(dim) - estimated_dim += 1 advanced_index[estimated_dim] = ( estimated_dim, - paddle.to_tensor(slice_item), + paddle.to_tensor([slice_item]), ) has_advanced_index = True + estimated_dim += 1 elif isinstance(slice_item, slice): start = slice_item.start end = slice_item.stop step = slice_item.step estimated_dim += 1 - + dim += 1 if start is None and end is None and step is None: continue @@ -760,7 +762,7 @@ def parse_index(x, indices): has_advanced_index = True estimated_dim += 1 - + dim += 1 elif isinstance(slice_item, paddle.base.Variable): # In this case, the Variable is not 0-dim Tensor and will be treated as advanced-indexing. if ( @@ -780,7 +782,7 @@ def parse_index(x, indices): advanced_index[estimated_dim] = (estimated_dim, slice_item) has_advanced_index = True estimated_dim += 1 - + dim += 1 elif isinstance(slice_item, paddle.pir.OpResult): # In this case, the Variable is not 0-dim Tensor and will be treated as advanced-indexing. if slice_item.dtype == paddle.pir.core.DataType.BOOL: @@ -797,7 +799,7 @@ def parse_index(x, indices): advanced_index[estimated_dim] = (estimated_dim, slice_item) has_advanced_index = True estimated_dim += 1 - + dim += 1 else: raise IndexError( "Valid index accept int / bool / slice / ellipsis / list / Tuple / Ndarray / Tensor, but received {}.".format( @@ -808,7 +810,7 @@ def parse_index(x, indices): starts.append(start) ends.append(end) steps.append(step) - axes.append(dim) + axes.append(dim - 1) use_strided_slice = ( True if ( diff --git a/test/indexing/test_getitem.py b/test/indexing/test_getitem.py index 6ce8519926a49..f3a2374ecbe1d 100644 --- a/test/indexing/test_getitem.py +++ b/test/indexing/test_getitem.py @@ -34,7 +34,7 @@ def test_combined_index_1(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[[0, 1], :, [1, 2]] @@ -52,7 +52,7 @@ def test_combined_index_2(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -70,7 +70,7 @@ def test_combined_index_3(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[[1, 0], :, [1, 4], 1:5:2, 4] @@ -89,7 +89,7 @@ def test_combined_index_4(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[:, [1, 0], 0:4:2, [2, 3], 4] @@ -107,7 +107,7 @@ def test_combined_index_5(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[::2, [1, 0], [2, 3], 0:4:2] @@ -125,7 +125,7 @@ def test_combined_index_6(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[::2, [1, 0], [2, 3], 0:4:2, [4, 6]] @@ -143,7 +143,7 @@ def test_combined_index_7(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[::2, [[1, 0]], [[2, 3]], 0:4:2, [[4, 6]]] @@ -161,7 +161,7 @@ def test_combined_index_8(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[ @@ -181,7 +181,7 @@ def test_combined_index_9(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[[[1, 0]], [1, 0], 0:4:2, [[3, 5], [4, 2]]] @@ -199,7 +199,7 @@ def test_combined_index_10(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[:, [True, False, True, False], 4] @@ -220,7 +220,7 @@ def test_combined_index_11(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[:, [False, False, False, False], 4] @@ -240,7 +240,7 @@ def test_index_has_range(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[:, range(3), 4] @@ -261,7 +261,7 @@ def test_indexing_with_bool_list1(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[[True, False, True], [False, False, False, True]] @@ -282,7 +282,7 @@ def test_indexing_with_bool_list2(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[ @@ -311,7 +311,7 @@ def test_indexing_is_multi_dim_list(self): if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_res = np_data[np.array([[2, 3, 4], [1, 2, 5]])] @@ -326,6 +326,45 @@ def test_indexing_is_multi_dim_list(self): np.testing.assert_allclose(y.numpy(), np_res) np.testing.assert_allclose(y.numpy(), y_index_tensor.numpy()) + def test_indexing_is_boolean_true(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = ( + np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)).astype(self.ndtype) + ) + + if self.dtype == 'bfloat16': + np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) + if self.dtype == 'complex64' or self.dtype == 'complex128': + np_data = np_data + 1j * np_data + + np_res = np_data[True] + + x = paddle.to_tensor(np_data, dtype=self.dtype) + y = x[True] + + if self.dtype == 'bfloat16': + y = paddle.cast(y, dtype='float32') + + np.testing.assert_allclose(y.numpy(), np_res) + + def test_indexing_is_boolean_false(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = ( + np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)).astype(self.ndtype) + ) + + if self.dtype == 'bfloat16': + np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) + if self.dtype == 'complex64' or self.dtype == 'complex128': + np_data = np_data + 1j * np_data + + np_res = np_data[1, False, 0] + + x = paddle.to_tensor(np_data, dtype=self.dtype) + y = x[1, False, 0] + + np.testing.assert_allclose(y.numpy(), np_res) + @unittest.skipIf( not core.is_compiled_with_cuda() @@ -1002,6 +1041,36 @@ def test_indexing_is_multi_dim_list(self): np.testing.assert_allclose(res[0], np_res) np.testing.assert_allclose(res[1], np_res) + def test_indexing_is_boolean_true(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)) + np_res = np_data[True] + + with paddle.static.program_guard( + paddle.static.Program(), paddle.static.Program() + ): + x = paddle.to_tensor(np_data) + y = _getitem_static(x, True) + + res = self.exe.run(fetch_list=[y.name]) + + np.testing.assert_allclose(res[0], np_res) + + def test_indexing_is_boolean_false(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)) + np_res = np_data[1, False, 0] + + with paddle.static.program_guard( + paddle.static.Program(), paddle.static.Program() + ): + x = paddle.to_tensor(np_data) + y = _getitem_static(x, (1, False, 0)) + + res = self.exe.run(fetch_list=[y.name]) + + np.testing.assert_allclose(res[0], np_res) + class TestGetitemBasicIndexOutputView(unittest.TestCase): def setUp(self): diff --git a/test/indexing/test_setitem.py b/test/indexing/test_setitem.py index 9a39b4b56d221..b8d7e3361efc4 100644 --- a/test/indexing/test_setitem.py +++ b/test/indexing/test_setitem.py @@ -32,7 +32,7 @@ def test_combined_index_1(self): np_data = np.zeros((3, 4, 5, 6), dtype='float32').astype(self.ndtype) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -47,7 +47,7 @@ def test_combined_index_2(self): np_data = np.ones((3, 4, 5, 6), dtype='float32').astype(self.ndtype) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -62,7 +62,7 @@ def test_combined_index_3(self): np_data = np.ones((3, 4, 5, 6), dtype='int32').astype(self.ndtype) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -77,7 +77,7 @@ def test_index_has_range(self): np_data = np.ones((3, 4, 5, 6), dtype='int32').astype(self.ndtype) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -93,7 +93,7 @@ def test_src_value_with_different_dtype_1(self): np_data = np.ones((3, 4, 5, 6), dtype='int32').astype(self.ndtype) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_value = np.zeros((6,), dtype='float32') x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -111,7 +111,7 @@ def test_src_value_with_different_dtype_2(self): np_data = np.ones((3, 4, 5, 6), dtype='float32').astype(self.ndtype) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data np_value = np.zeros((6,), dtype='int64') @@ -132,7 +132,7 @@ def test_indexing_with_bool_list1(self): ) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -151,7 +151,7 @@ def test_indexing_with_bool_list2(self): ) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) @@ -178,7 +178,7 @@ def test_indexing_is_multi_dim_list(self): ) if self.dtype == 'bfloat16': np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_data = np_data + 1j * np_data x = paddle.to_tensor(np_data, dtype=self.dtype) np_data[np.array([[2, 3, 4], [1, 2, 5]])] = 100 @@ -188,11 +188,51 @@ def test_indexing_is_multi_dim_list(self): x = paddle.cast(x, dtype='float32') np.testing.assert_allclose(x.numpy(), np_data) + def test_indexing_is_boolean_true(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = ( + np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)).astype(self.ndtype) + ) + + if self.dtype == 'bfloat16': + np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) + if self.dtype == 'complex64' or self.dtype == 'complex128': + np_data = np_data + 1j * np_data + + x = paddle.to_tensor(np_data, dtype=self.dtype) + np_data[2, True, :, 1] = 100 + x[2, True, :, 1] = 100 + + if self.dtype == 'bfloat16': + x = paddle.cast(x, dtype='float32') + + np.testing.assert_allclose(x.numpy(), np_data) + + def test_indexing_is_boolean_false(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = ( + np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)).astype(self.ndtype) + ) + + if self.dtype == 'bfloat16': + np_data = convert_uint16_to_float(convert_float_to_uint16(np_data)) + if self.dtype == 'complex64' or self.dtype == 'complex128': + np_data = np_data + 1j * np_data + + x = paddle.to_tensor(np_data, dtype=self.dtype) + np_data[False] = 100 + x[False] = 100 + + if self.dtype == 'bfloat16': + x = paddle.cast(x, dtype='float32') + + np.testing.assert_allclose(x.numpy(), np_data) + def test_inplace_with_stride(self): np_v = np.random.randn(3, 1).astype(self.ndtype) if self.dtype == 'bfloat16': np_v = convert_uint16_to_float(convert_float_to_uint16(np_v)) - if self.dtype == 'complex64' or self.dtype == 'complex32': + if self.dtype == 'complex64' or self.dtype == 'complex128': np_v = np_v + 1j * np_v v = paddle.to_tensor(np_v, dtype=self.dtype) v.stop_gradient = False @@ -504,6 +544,36 @@ def test_indexing_is_multi_dim_list(self): np.testing.assert_allclose(res[0], np_data) + def test_indexing_is_boolean_true(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)) + np_data[2, True, :, 1] = 100 + + with paddle.static.program_guard( + paddle.static.Program(), paddle.static.Program() + ): + x = paddle.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)) + y = _setitem_static(x, (2, True, slice(None), 1), 100) + + res = self.exe.run(fetch_list=[y.name]) + + np.testing.assert_allclose(res[0], np_data) + + def test_indexing_is_boolean_false(self): + # indexing is boolean, should improve rank of tensor and then treat it as advanced indexing. + np_data = np.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)) + np_data[False] = 100 + + with paddle.static.program_guard( + paddle.static.Program(), paddle.static.Program() + ): + x = paddle.arange(3 * 4 * 5 * 6).reshape((6, 5, 4, 3)) + y = _setitem_static(x, False, 100) + + res = self.exe.run(fetch_list=[y.name]) + + np.testing.assert_allclose(res[0], np_data) + if __name__ == '__main__': unittest.main()