add test file

python/paddle/__init__.py

@@ -258,6 +258,8 @@
     unfold,
     masked_fill,
     masked_fill_,
+    masked_scatter_,
+    masked_scatter,
     index_fill,
     index_fill_,
     diagonal_scatter,
@@ -926,6 +928,8 @@
     'polygamma_',
     'masked_fill',
     'masked_fill_',
+    'masked_scatter',
+    'masked_scatter_',
     'hypot',
     'hypot_',
     'index_fill',

python/paddle/tensor/__init__.py

@@ -171,6 +171,8 @@
 from .manipulation import unfold  # noqa: F401
 from .manipulation import masked_fill  # noqa: F401
 from .manipulation import masked_fill_  # noqa: F401
+from .manipulation import masked_scatter
+from .manipulation import masked_scatter_
 from .manipulation import index_fill  # noqa: F401
 from .manipulation import index_fill_  # noqa: F401
 from .manipulation import diagonal_scatter  # noqa: F401
@@ -743,6 +745,8 @@
     'atleast_2d',
     'atleast_3d',
     'diagonal_scatter',
+    'masked_scatter',
+    'masked_scatter_',
 ]
 
 # this list used in math_op_patch.py for magic_method bind

python/paddle/tensor/manipulation.py

@@ -4009,9 +4009,37 @@ def get_attr_shape(list_shape):
         return out
 
 
-def masked_scatter(x, mask, value):
+def masked_scatter(x, mask, value, name=None):
     """
-    利用现有api实现masked_scatter功能
+    Copies elements from `value` into `x` tensor at positions where the `mask` is True.
+
+    Elements from source are copied into `x` starting at position 0 of `value` and continuing in order one-by-one for
+    each occurrence of `mask` being True. The shape of `mask` must be broadcastable with the shape of the underlying tensor.
+    The `value` should have at least as many elements as the number of ones in `mask`.
+
+    Args:
+        x (Tensor): An N-D Tensor. The data type is ``float16``, ``float32``, ``float64``, ``int16``, ``int32``, ``int64``, ``int8``, ``uint8``, ``complex64``, ``complex128``, ``bfloat16`` or ``bool``.
+        shape (list|tuple|Tensor): Define the target shape. At most one dimension of the target shape can be -1.
+                        The data type is ``int32`` . If ``shape`` is a list or tuple, each element of it should be integer or Tensor with shape [].
+                        If ``shape`` is an Tensor, it should be an 1-D Tensor .
+        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        Tensor, A reshaped Tensor with the same data type as ``x``.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+
+            >>> x = paddle.randn([3,4])
+            >>> mask = paddle.randn([3,4])
+            >>> mask = mask>0.6
+            >>> value = paddle.ones([2,4], dtype="float32")
+
+            >>> out = paddle.masked_scatter(x, mask, value)
+            >>> print(out.shape)
+
     """
     # make sure the dtype of x and source is the same
     assert (
@@ -4031,19 +4059,16 @@ def masked_scatter(x, mask, value):
 
 
 @inplace_apis_in_dygraph_only
-def masked_scatter_(x, mask, value):
+def masked_scatter_(x, mask, value, name=None):
     """
-    利用现有api实现masked_scatter功能
+    Inplace version of ``masked_scatter`` API, the output Tensor will be inplaced with input ``x``.
+    Please refer to :ref:`api_paddle_masked_scatter`.
     """
     # make sure the dtype of x and source is the same
     assert (
         x.dtype == value.dtype
     ), f'x and value must have the same dtype, but got x dtype is {x.dtype}, value dtype is {value.dtype}'
     assert mask.dtype == paddle.bool
-    if not x.stop_gradient:
-        raise ValueError(
-            "The inplace operation only support stop_gradient=True inplace tensor"
-        )
     shape = paddle.broadcast_shape(x.shape, mask.shape)
     if shape != x.shape:
         raise ValueError(
@@ -4059,7 +4084,8 @@ def masked_scatter_(x, mask, value):
 
     value = value.flatten()[mask_prefix].reshape(mask.shape)
     mask = paddle.logical_not(mask)
-    paddle.where_(mask, x, value)
+    out = paddle.where_(mask, x, value)
+    return out
 
 
 @inplace_apis_in_dygraph_only

test/legacy_test/test_inplace.py

@@ -316,6 +316,22 @@ def init_data(self):
         self.mask = paddle.to_tensor(self.mask, dtype='bool')
 
 
+class TestDygraphInplaceMaskedScatter(TestDygraphInplace):
+    def non_inplace_api_processing(self, var):
+        return paddle.masked_scatter(var, self.mask, self.value)
+
+    def inplace_api_processing(self, var):
+        return paddle.masked_scatter_(var, self.mask, self.value)
+
+    def init_data(self):
+        self.dtype = "float32"
+        self.input_var_numpy = np.random.uniform(-5, 5, [30, 3])
+        self.value = np.random.uniform(size=(30, 30))
+        self.value = paddle.to_tensor(self.value, dtype=self.dtype)
+        self.mask = np.random.randint(0, 2, [30, 1]).astype('bool')
+        self.mask = paddle.to_tensor(self.mask, dtype='bool')
+
+
 class TestDygraphInplaceWithContinuous(TestDygraphInplace):
     def init_data(self):
         self.input_var_numpy = np.random.uniform(-5, 5, [10, 20, 1])