[Paddle Tensor No.23] 新增TensorA[mask] = TensorB (#69480)

LittleHeroZZZX · web-flow · commit 6c3c4fac0954 · 2024-11-20T16:28:09.000+08:00
* Fix Tensor.__setitem__ to correctly support boolean mask assignment

* retain broadcast-then-where logic for scalar values in setitem
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1863,13 +1863,11 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
         ConvertAllInputsToDistTensor(
             mesh, self->tensor, transed_sub_tensor, value_tensor);
       }
-
       if (transed_index.size() == 1 &&
           transed_index[0].dtype() == phi::DataType::BOOL &&
-          transed_index[0].shape().size() == self->tensor.shape().size()) {
-        if (value_tensor.shape() != self->tensor.shape()) {
-          value_tensor = expand_ad_func(value_tensor, self->tensor.shape());
-        }
+          transed_index[0].shape().size() == self->tensor.shape().size() &&
+          value_tensor.numel() == 1) {
+        value_tensor = expand_ad_func(value_tensor, self->tensor.shape());
         transed_sub_tensor =
             where__ad_func(logical_not_ad_func(transed_index[0]),
                            transed_sub_tensor,
@@ -1878,7 +1876,6 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
         transed_sub_tensor =
             index_put__ad_func(transed_sub_tensor, transed_index, value_tensor);
       }
-
       if (out_is_view) {
         // NOTE(zoooo0820): if out_is_view is true, it is a case of
         // combined-indexing setitem, i.e. firstly we get a view of
diff --git a/test/indexing/test_setitem.py b/test/indexing/test_setitem.py
@@ -533,6 +533,28 @@ def test_basic_setitem_bwd_2(self):
         else:
             np.testing.assert_equal(v.grad.numpy(), expected_v_grad)
 
+    def test_boolean_mask_scalar(self):
+        tensor_np = np.arange(2 * 3).reshape(2, 3)
+        tensor = paddle.to_tensor(tensor_np)
+        mask_np = np.array([[True, True, False], [False, True, False]])
+        mask = paddle.to_tensor(mask_np)
+        tensor[mask] = 100
+        tensor_np[mask_np] = 100
+        np.testing.assert_equal(tensor.numpy(), tensor_np)
+
+    def test_boolean_mask_tensor(self):
+        tensor_np = np.arange(2 * 3).reshape(2, 3)
+        tensor = paddle.to_tensor(tensor_np)
+        mask_np = np.array([[True, True, False], [False, True, False]]).astype(
+            'bool'
+        )
+        mask = paddle.to_tensor(mask_np)
+        value_np = np.array([100] * mask_np.sum())
+        value = paddle.to_tensor(value_np)
+        tensor[mask] = value
+        tensor_np[mask_np] = value_np
+        np.testing.assert_equal(tensor.numpy(), tensor_np)
+
 
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
@@ -941,6 +963,38 @@ def test_combined_indexing_and_value_is_tensor_3(self):
 
         np.testing.assert_allclose(res[0], np_data)
 
+    def test_boolean_mask_scalar(self):
+        tensor_np = np.arange(2 * 3).reshape(2, 3).astype('int32')
+        mask_np = np.array([[True, True, False], [False, True, False]]).astype(
+            'bool'
+        )
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            tensor = paddle.to_tensor(tensor_np)
+            mask = paddle.to_tensor(mask_np)
+            tensor[mask] = 100
+            res = self.exe.run(fetch_list=[tensor])
+        tensor_np[mask_np] = 100
+        np.testing.assert_equal(res[0], tensor_np)
+
+    def test_boolean_mask_tensor(self):
+        tensor_np = np.arange(2 * 3).reshape(2, 3).astype('int32')
+        mask_np = np.array([[True, True, False], [False, True, False]]).astype(
+            'bool'
+        )
+        value_np = np.array([100] * mask_np.sum()).astype('int32')
+        with paddle.static.program_guard(
+            paddle.static.Program(), paddle.static.Program()
+        ):
+            tensor = paddle.to_tensor(tensor_np)
+            mask = paddle.to_tensor(mask_np)
+            value = paddle.to_tensor(value_np)
+            tensor[mask] = value
+            res = self.exe.run(fetch_list=[tensor])
+        tensor_np[mask_np] = value_np
+        np.testing.assert_equal(res[0], tensor_np)
+
 
 if __name__ == '__main__':
     unittest.main()
