Fix shape error in combined-indexing setitem (#60447)

zoooo0820 · web-flow · commit 1761931b9c30 · 2024-01-02T12:07:21.000+08:00
* add ut

* fix shape error in combine-indexing

* fix ut
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1375,7 +1375,7 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
 
   // step3: Dealing with advanced indexing
   std::vector<paddle::Tensor> transed_index;
-  std::vector<int> trans_back_dim;
+  std::vector<int> trans_back_dim, trans_dim;
   int pos_of_new_dim = INT_MAX, rank_of_new_dim = 1;
 
   paddle::Tensor transed_tensor = dealWithAdvancedIndex(out,
@@ -1385,7 +1385,8 @@ static PyObject* tensor__getitem_dygraph(TensorObject* self,
                                                         &transed_index,
                                                         &trans_back_dim,
                                                         &pos_of_new_dim,
-                                                        &rank_of_new_dim);
+                                                        &rank_of_new_dim,
+                                                        &trans_dim);
 
   if (transed_index.size() == 1 &&
       transed_index[0].dtype() == phi::DataType::BOOL) {
@@ -1679,9 +1680,9 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
                                                            &use_strided_slice);
 
     std::vector<paddle::Tensor> transed_index;
-    std::vector<int> trans_back_dim;
+    std::vector<int> trans_back_dim, trans_dim;
 
-    int pos_of_new_dim = 0, rank_of_new_dim = 0;
+    int pos_of_new_dim = INT_MAX, rank_of_new_dim = 1;
 
     paddle::Tensor transed_sub_tensor =
         dealWithAdvancedIndex(sub_tensor,
@@ -1691,7 +1692,8 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
                               &transed_index,
                               &trans_back_dim,
                               &pos_of_new_dim,
-                              &rank_of_new_dim);
+                              &rank_of_new_dim,
+                              &trans_dim);
 
     // Release gil and do tracing
     py::gil_scoped_release release;
@@ -1714,6 +1716,10 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
       }
     }
 
+    if (value_tensor.dims().size() > 1 && pos_of_new_dim != 0) {
+      value_tensor = transpose_ad_func(value_tensor, trans_dim);
+    }
+
     // TODO(zoooo0820) 1.Using inplace version index_put
     //                  2.Remove following code after backward bug fixed.
     transed_sub_tensor = assign_ad_func(transed_sub_tensor);
diff --git a/paddle/fluid/pybind/slice_utils.h b/paddle/fluid/pybind/slice_utils.h
@@ -398,9 +398,8 @@ static paddle::Tensor dealWithAdvancedIndex(
     std::vector<paddle::Tensor>* transed_index,
     std::vector<int>* trans_back_dim,
     int* pos_of_new_dim,
-    int* rank_of_new_dim) {
-  std::vector<int> trans_dim;
-
+    int* rank_of_new_dim,
+    std::vector<int>* trans_dim) {
   int p = 0;
   for (size_t i = 0; i < advanced_index_dim->size(); ++i) {
     auto index_dim = (*advanced_index_dim)[i];
@@ -409,30 +408,28 @@ static paddle::Tensor dealWithAdvancedIndex(
       // advanced_index_dim
       auto index = (*advanced_index)[p++];
 
-      if (!is_for_setitem) {
-        if (index_dim == 0) {
-          // case 1: advanced indices at axis 0, the new dim will be at first.
-          *pos_of_new_dim = 0;
-        } else if (index_dim > 0 && trans_dim.size() > 0 &&
-                   trans_dim[trans_dim.size() - 1] != index_dim - 1) {
-          // case 2: there are not adjacent advanced indices, the new dim will
-          // be at first.
-          *pos_of_new_dim = 0;
-        } else {
-          *pos_of_new_dim = std::min(index_dim, *pos_of_new_dim);
-        }
-        *rank_of_new_dim =
-            std::max(*rank_of_new_dim, static_cast<int>(index.shape().size()));
+      if (index_dim == 0) {
+        // case 1: advanced indices at axis 0, the new dim will be at first.
+        *pos_of_new_dim = 0;
+      } else if (index_dim > 0 && trans_dim->size() > 0 &&
+                 (*trans_dim)[trans_dim->size() - 1] != index_dim - 1) {
+        // case 2: there are not adjacent advanced indices, the new dim will
+        // be at first.
+        *pos_of_new_dim = 0;
+      } else {
+        *pos_of_new_dim = std::min(index_dim, *pos_of_new_dim);
       }
+      *rank_of_new_dim =
+          std::max(*rank_of_new_dim, static_cast<int>(index.shape().size()));
 
-      trans_dim.push_back(index_dim);
+      trans_dim->push_back(index_dim);
       transed_index->push_back(std::move(index));
     }
   }
 
   for (size_t i = 0; i < tensor.shape().size(); ++i) {
     if ((*advanced_index_dim)[i] == -1) {
-      trans_dim.push_back(i);
+      trans_dim->push_back(i);
     }
   }
 
@@ -442,19 +439,19 @@ static paddle::Tensor dealWithAdvancedIndex(
   std::vector<int> original_dim_order(tensor.shape().size());
   std::iota(original_dim_order.begin(), original_dim_order.end(), 0);
 
-  if (original_dim_order == trans_dim) {
+  if (original_dim_order == *trans_dim) {
     transed_tensor = tensor;
   } else {
-    transed_tensor = transpose_ad_func(tensor, trans_dim);
+    transed_tensor = transpose_ad_func(tensor, *trans_dim);
   }
 
   if (is_for_setitem) {
-    trans_back_dim->resize(trans_dim.size());
+    trans_back_dim->resize(trans_dim->size());
     std::iota(trans_back_dim->begin(), trans_back_dim->end(), 0);
     std::sort(trans_back_dim->begin(),
               trans_back_dim->end(),
               [&trans_dim](int left, int right) {
-                return trans_dim[left] < trans_dim[right];
+                return (*trans_dim)[left] < (*trans_dim)[right];
               });
   }
   return transed_tensor;
diff --git a/python/paddle/base/variable_index.py b/python/paddle/base/variable_index.py
@@ -170,7 +170,7 @@ def _setitem_for_tensor_array(var, item, value):
         )
 
 
-def deal_advanced_index(ori_tensor, indices, is_for_setitem):
+def deal_advanced_index(ori_tensor, indices, is_for_setitem, values):
     """
     Transpose origin Tensor and advanced indices to the front.
 
@@ -180,6 +180,7 @@ def deal_advanced_index(ori_tensor, indices, is_for_setitem):
         trans_back_dim (List): order of axes to transpose back to original order. Only used in __setitem__.
         pos_of_new_dim (int):  axis of new dim in the result. Only used in __getitem__.
         rank_of_new_dim (int): rank of new dim in the result. Only used in __getitem__.
+        transed_value_tensor (Tensor): value tensor transed to the front. Only used in __setitem__.
     """
     transed_dim = []
     transed_index = []
@@ -191,16 +192,15 @@ def deal_advanced_index(ori_tensor, indices, is_for_setitem):
 
     for i, indice in enumerate(indices):
         if indice is not None:
-            if not is_for_setitem:
-                if i == 0:
-                    # case 1: advanced indices at axis 0, the new dim will be at first.
-                    pos_of_new_dim = 0
-                if i > 0 and len(transed_dim) > 0 and transed_dim[-1] != i - 1:
-                    # case 2: there are not adjacent advanced indices, the new dim will be at first.
-                    pos_of_new_dim = 0
-                else:
-                    pos_of_new_dim = min(pos_of_new_dim, i)
-                rank_of_new_dim = max(rank_of_new_dim, indice[1].ndim)
+            if i == 0:
+                # case 1: advanced indices at axis 0, the new dim will be at first.
+                pos_of_new_dim = 0
+            if i > 0 and len(transed_dim) > 0 and transed_dim[-1] != i - 1:
+                # case 2: there are not adjacent advanced indices, the new dim will be at first.
+                pos_of_new_dim = 0
+            else:
+                pos_of_new_dim = min(pos_of_new_dim, i)
+            rank_of_new_dim = max(rank_of_new_dim, indice[1].ndim)
             transed_dim.append(i)
             transed_index.append(indice[1])
     for i in range(ori_tensor.ndim):
@@ -210,12 +210,22 @@ def deal_advanced_index(ori_tensor, indices, is_for_setitem):
 
     trans_back_dim = np.argsort(transed_dim).tolist() if is_for_setitem else []
 
+    transed_value_tensor = None
+    if is_for_setitem:
+        if values.ndim > 1 and pos_of_new_dim != 0:
+            # If the value tensor is not a scalar / 1-D Tensor, and the src tensor was
+            # transposed at 1st dim, the value tensor should be transposed too.
+            transed_value_tensor = values.transpose(transed_dim)
+        else:
+            transed_value_tensor = values
+
     return (
         transed_tensor,
         transed_index,
         trans_back_dim,
         pos_of_new_dim,
         rank_of_new_dim,
+        transed_value_tensor,
     )
 
 
@@ -577,6 +587,11 @@ def _setitem_static(x, indices, values):
         #   3. assign values to the sliced result by index_put OP;
         #   4. transpose back and assign the result to original tensor by set_value OP.
 
+        if not isinstance(
+            values, (Variable, paddle.pir.Value, paddle.pir.OpResult)
+        ):
+            values = paddle.assign(values).astype(x.dtype)
+
         sub_tensor = get_tensor_with_basic_indexing(
             x,
             axes,
@@ -593,9 +608,8 @@ def _setitem_static(x, indices, values):
             transback_dim,
             _,
             _,
-        ) = deal_advanced_index(sub_tensor, advanced_index, True)
-        if not isinstance(values, (Variable, paddle.pir.Value)):
-            values = paddle.assign(values).astype(transed_sub_tensor.dtype)
+            values,
+        ) = deal_advanced_index(sub_tensor, advanced_index, True, values)
 
         if values.dtype != transed_sub_tensor.dtype:
             values = values.astype(transed_sub_tensor.dtype)
@@ -818,7 +832,8 @@ def _getitem_static(x, indices):
             _,
             pos_of_new_dim,
             rank_of_new_dim,
-        ) = deal_advanced_index(out, advanced_index, False)
+            _,
+        ) = deal_advanced_index(out, advanced_index, False, None)
 
         # TODO(zooooo0820): Replacing gather_nd to another advanded OP for handling of mixed indexes more efficiently
         if (
diff --git a/test/indexing/test_setitem.py b/test/indexing/test_setitem.py
@@ -229,6 +229,79 @@ def test_indexing_is_boolean_false(self):
 
         np.testing.assert_allclose(x.numpy(), np_data)
 
+    def test_combined_indexing_and_value_is_tensor_1(self):
+        # value is tensor with same shape to getitem and index will be adjusted
+        np_data = np.ones((3, 3)).astype(self.ndtype)
+        value_data = np.array([-1, -1, -1]).astype(self.ndtype)
+
+        if self.dtype == 'bfloat16':
+            np_data = convert_uint16_to_float(convert_float_to_uint16(np_data))
+            value_data = convert_uint16_to_float(
+                convert_float_to_uint16(value_data)
+            )
+        if self.dtype == 'complex64' or self.dtype == 'complex128':
+            np_data = np_data + 1j * np_data
+            value_data = value_data + 1j * value_data
+
+        x = paddle.to_tensor(np_data, dtype=self.dtype)
+        v = paddle.to_tensor(value_data, dtype=self.dtype)
+
+        np_data[:, [0, 2]] = np_data[:, [0, 2]] + np.expand_dims(value_data, -1)
+        x[:, [0, 2]] = x[:, [0, 2]] + v.unsqueeze(-1)
+
+        if self.dtype == 'bfloat16':
+            x = paddle.cast(x, dtype='float32')
+
+        np.testing.assert_allclose(x.numpy(), np_data)
+
+    def test_combined_indexing_and_value_is_tensor_2(self):
+        # value is tensor needed to broadcast and index will be adjusted
+        np_data = np.ones((3, 4, 5, 6)).astype(self.ndtype)
+        value_data = np.arange(3 * 4 * 2 * 1).reshape((3, 4, 2, 1))
+
+        if self.dtype == 'bfloat16':
+            np_data = convert_uint16_to_float(convert_float_to_uint16(np_data))
+            value_data = convert_uint16_to_float(
+                convert_float_to_uint16(value_data)
+            )
+        if self.dtype == 'complex64' or self.dtype == 'complex128':
+            np_data = np_data + 1j * np_data
+            value_data = value_data + 1j * value_data
+
+        x = paddle.to_tensor(np_data, dtype=self.dtype)
+        v = paddle.to_tensor(value_data, dtype=self.dtype)
+        x[..., [1, 4], ::2] = v
+
+        np_data[..., [1, 4], ::2] = value_data
+        if self.dtype == 'bfloat16':
+            x = paddle.cast(x, dtype='float32')
+        np.testing.assert_allclose(x.numpy(), np_data)
+
+    def test_combined_indexing_and_value_is_tensor_3(self):
+        # value is tensor and index will be adjusted
+        # and the value rank is less than original tensor
+        np_data = np.ones((3, 4, 5, 6)).astype(self.ndtype)
+        value_data = np.arange(2 * 3 * 5).reshape((2, 3, 5))
+
+        if self.dtype == 'bfloat16':
+            np_data = convert_uint16_to_float(convert_float_to_uint16(np_data))
+            value_data = convert_uint16_to_float(
+                convert_float_to_uint16(value_data)
+            )
+        if self.dtype == 'complex64' or self.dtype == 'complex128':
+            np_data = np_data + 1j * np_data
+            value_data = value_data + 1j * value_data
+
+        x = paddle.to_tensor(np_data, dtype=self.dtype)
+        v = paddle.to_tensor(value_data, dtype=self.dtype)
+        x[:, [1, 3], :, [3, 4]] = v
+
+        np_data[:, [1, 3], :, [3, 4]] = value_data
+
+        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':
@@ -588,6 +661,72 @@ def test_indexing_is_boolean_false(self):
 
         np.testing.assert_allclose(res[0], np_data)
 
+    @test_with_pir_api
+    def test_combined_indexing_and_value_is_tensor_1(self):
+        # value is tensor with same shape to getitem and index will be adjusted
+        np_data = np.ones((3, 3), dtype='int32')
+        value_data = np.array([-1, -1, -1])
+        np_data[:, [0, 2]] = np_data[:, [0, 2]] * np.expand_dims(value_data, -1)
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            x = paddle.ones((3, 3), dtype='int32')
+            v = paddle.to_tensor([-1, -1, -1])
+            y = _setitem_static(
+                x,
+                (slice(None), [0, 2]),
+                x[:, [0, 2]] * v.unsqueeze(-1),
+            )
+            res = self.exe.run(fetch_list=[y])
+
+        np.testing.assert_allclose(res[0], np_data)
+
+    @test_with_pir_api
+    def test_combined_indexing_and_value_is_tensor_2(self):
+        # value is tensor needed to broadcast and index will be adjusted
+        np_data = np.ones((3, 4, 5, 6), dtype='int32')
+        value_data = np.arange(3 * 4 * 2 * 1).reshape((3, 4, 2, 1))
+        np_data[..., [1, 4], ::2] = value_data
+
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            x = paddle.ones((3, 4, 5, 6), dtype='int32')
+            v = paddle.arange(3 * 4 * 2 * 1).reshape((3, 4, 2, 1))
+
+            y = _setitem_static(
+                x,
+                (..., [1, 4], slice(None, None, 2)),
+                v,
+            )
+
+            res = self.exe.run(fetch_list=[y])
+
+        np.testing.assert_allclose(res[0], np_data)
+
+    @test_with_pir_api
+    def test_combined_indexing_and_value_is_tensor_3(self):
+        # value is tensor and index will be adjusted
+        # and the value rank is less than original tensor
+        np_data = np.ones((3, 4, 5, 6), dtype='int32')
+        value_data = np.arange(2 * 3 * 5).reshape((2, 3, 5))
+        np_data[:, [1, 3], :, [3, 4]] = value_data
+
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            x = paddle.ones((3, 4, 5, 6), dtype='int32')
+            v = paddle.arange(2 * 3 * 5).reshape((2, 3, 5))
+            y = _setitem_static(
+                x,
+                (slice(None), [1, 3], slice(None), [3, 4]),
+                v,
+            )
+
+            res = self.exe.run(fetch_list=[y])
+
+        np.testing.assert_allclose(res[0], np_data)
+
 
 if __name__ == '__main__':
     unittest.main()
