【Hackathon 6th No.11】为 Paddle 新增 bernoulli_ API - part (#64252)

NKNaN · web-flow · commit fba46752c14a · 2024-05-21T17:17:19.000+08:00
* add bernoulli_

* update

* update docs

* fix typo

* update code and test

* add test case for backward

* add test broadcast error and update docs
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -516,6 +516,7 @@
 )
 from .tensor.random import (
     bernoulli,
+    bernoulli_,
     binomial,
     check_shape,
     multinomial,
@@ -808,6 +809,7 @@
     'expm1',
     'expm1_',
     'bernoulli',
+    'bernoulli_',
     'binomial',
     'poisson',
     'standard_gamma',
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
@@ -402,6 +402,7 @@
     vander,
 )
 from .random import (  # noqa: F401
+    bernoulli_,
     binomial,
     exponential_,
     multinomial,
@@ -751,6 +752,7 @@
     'put_along_axis',
     'select_scatter',
     'put_along_axis_',
+    'bernoulli_',
     'exponential_',
     'heaviside',
     'index_add',
diff --git a/python/paddle/tensor/random.py b/python/paddle/tensor/random.py
@@ -104,6 +104,58 @@ def bernoulli(x, name=None):
         return out
 
 
+@dygraph_only
+def bernoulli_(x, p=0.5, name=None):
+    """
+    This is the inplace version of api ``bernoulli``, which returns a Tensor filled
+    with random values sampled from a bernoulli distribution. The output Tensor will
+    be inplaced with input ``x``. Please refer to :ref:`api_tensor_bernoulli`.
+
+    Args:
+        x(Tensor): The input tensor to be filled with random values.
+        p (float|Tensor, optional): The success probability parameter of the output Tensor's bernoulli distribution.
+            If ``p`` is float, all elements of the output Tensor shared the same success probability.
+            If ``p`` is a Tensor, it has per-element success probabilities, and the shape should be broadcastable to ``x``.
+            Default is 0.5
+        name(str, optional): The default value is None. Normally there is no
+            need for user to set this property. For more information, please
+            refer to :ref:`api_guide_Name`.
+
+    Returns:
+        A Tensor filled with random values sampled from the bernoulli distribution with success probability ``p`` .
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> x = paddle.randn([3, 4])
+            >>> x.bernoulli_()
+            >>> # doctest: +SKIP('random check')
+            >>> print(x)
+            Tensor(shape=[3, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[1., 0., 1., 0.],
+             [0., 1., 1., 0.],
+             [0., 1., 1., 1.]])
+
+            >>> x = paddle.randn([3, 4])
+            >>> p = paddle.randn([3, 1])
+            >>> x.bernoulli_(p)
+            >>> # doctest: +SKIP('random check')
+            >>> print(x)
+            Tensor(shape=[3, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[1., 0., 0., 1.],
+             [1., 0., 1., 0.],
+             [1., 0., 0., 0.]])
+
+    """
+    x.uniform_(0.0, 1.0)
+    ones_mask = x > p
+    zeros_mask = x < p
+    x.masked_fill_(ones_mask, 1.0)
+    x.masked_fill_(zeros_mask, 0.0)
+    return x
+
+
 def binomial(count, prob, name=None):
     r"""
     Returns a tensor filled with random number from the Binomial Distribution, which supports Tensor shape
diff --git a/test/deprecated/legacy_test/test_inplace.py b/test/deprecated/legacy_test/test_inplace.py
@@ -1921,5 +1921,97 @@ def test_inplace_api(self):
         )
 
 
+class TestDygraphInplaceBernoulli(unittest.TestCase):
+    def setUp(self):
+        self.init_data()
+        self.set_np_compare_func()
+
+    def init_data(self):
+        self.shape = (100, 1000)
+        self.input_var_numpy = np.random.random(self.shape)
+        self.dtype = "float32"
+        self.p = 0.5
+
+    def set_np_compare_func(self):
+        self.np_compare = np.array_equal
+
+    def inplace_api_processing(self, var):
+        return paddle.bernoulli_(var, p=self.p)
+
+    def inplace_class_method_processing(self, var):
+        return var.bernoulli_(self.p)
+
+    def non_inplace_api_processing(self):
+        return paddle.bernoulli(paddle.full(self.shape, self.p))
+
+    def test_inplace_api(self):
+        var = paddle.to_tensor(self.input_var_numpy).astype(self.dtype)
+        non_inplace_var = self.non_inplace_api_processing()
+        inplace_var = self.inplace_api_processing(var)
+        self.assertTrue(id(var) == id(inplace_var))
+        np.testing.assert_allclose(
+            non_inplace_var.numpy().mean(),
+            inplace_var.numpy().mean(),
+            atol=0.01,
+        )
+        np.testing.assert_allclose(
+            non_inplace_var.numpy().var(), inplace_var.numpy().var(), atol=0.01
+        )
+
+    def test_inplace_api_backward(self):
+        var_a = paddle.to_tensor(self.input_var_numpy).astype(self.dtype)
+        var_a.stop_gradient = False
+        var_b = var_a.clone()
+        expected_gradient = np.zeros(self.shape)
+        inplace_var = self.inplace_api_processing(var_b)
+        inplace_var.backward()
+        np.testing.assert_equal(
+            var_a.grad.numpy(),
+            expected_gradient,
+        )
+
+    def test_inplace_class_method(self):
+        var = paddle.to_tensor(self.input_var_numpy).astype(self.dtype)
+        non_inplace_var = self.non_inplace_api_processing()
+        inplace_var = self.inplace_class_method_processing(var)
+        self.assertTrue(id(var) == id(inplace_var))
+        np.testing.assert_allclose(
+            non_inplace_var.numpy().mean(),
+            inplace_var.numpy().mean(),
+            atol=0.01,
+        )
+        np.testing.assert_allclose(
+            non_inplace_var.numpy().var(), inplace_var.numpy().var(), atol=0.01
+        )
+
+    def test_inplace_class_method_backward(self):
+        var_a = paddle.to_tensor(self.input_var_numpy).astype(self.dtype)
+        var_a.stop_gradient = False
+        var_b = var_a.clone()
+        expected_gradient = np.zeros(self.shape)
+        inplace_var = self.inplace_class_method_processing(var_b)
+        inplace_var.backward()
+        np.testing.assert_equal(
+            var_a.grad.numpy(),
+            expected_gradient,
+        )
+
+
+class TestDygraphInplaceBernoulli2(TestDygraphInplaceBernoulli):
+    def init_data(self):
+        self.shape = (100, 1000)
+        self.input_var_numpy = np.random.random(self.shape)
+        self.dtype = "float64"
+        self.p = 0.5
+
+
+class TestDygraphInplaceBernoulliError(unittest.TestCase):
+    def test_broadcast_error(self):
+        var = paddle.randn([3, 4])
+        p = paddle.randn([5])
+        with self.assertRaises(ValueError):
+            var.bernoulli_(p)
+
+
 if __name__ == '__main__':
     unittest.main()
