From e9dd763cea08a1015fbb09b574efcddc118c2324 Mon Sep 17 00:00:00 2001
From: pangyoki <pangyoki@126.com>
Date: Wed, 16 Sep 2020 19:37:32 +0000
Subject: [PATCH] add multinomial python api

---
 python/paddle/__init__.py                     |  1 +
 .../tests/unittests/test_multinomial_op.py    | 20 +++++++++
 python/paddle/tensor/__init__.py              |  1 +
 python/paddle/tensor/random.py                | 42 +++++++++++++++++++
 4 files changed, 64 insertions(+)

diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
index 016726633ea..d1a42bef4a5 100755
--- a/python/paddle/__init__.py
+++ b/python/paddle/__init__.py
@@ -199,6 +199,7 @@ from .tensor.math import isfinite  #DEFINE_ALIAS
 from .tensor.math import isinf  #DEFINE_ALIAS
 from .tensor.math import isnan  #DEFINE_ALIAS
 from .tensor.math import prod  #DEFINE_ALIAS
+from .tensor.random import multinomial  #DEFINE_ALIAS
 from .tensor.random import standard_normal
 from .tensor.random import normal
 from .tensor.random import uniform  #DEFINE_ALIAS
diff --git a/python/paddle/fluid/tests/unittests/test_multinomial_op.py b/python/paddle/fluid/tests/unittests/test_multinomial_op.py
index bb78dd4d007..d4f9a176831 100644
--- a/python/paddle/fluid/tests/unittests/test_multinomial_op.py
+++ b/python/paddle/fluid/tests/unittests/test_multinomial_op.py
@@ -102,5 +102,25 @@ class TestReplacementError(unittest.TestCase):
         self.attrs = {"num_samples": 10, "replacement": False}
 """
 
+
+class TestMultinomialApi(unittest.TestCase):
+    def test_dygraph(self):
+        paddle.disable_static()
+        x = paddle.rand([4])
+        out = paddle.multinomial(x, num_samples=100000, replacement=True)
+        x_numpy = x.numpy()
+        paddle.enable_static()
+
+        sample_prob = np.unique(
+            out.numpy(), return_counts=True)[1].astype("float32")
+        sample_prob /= sample_prob.sum()
+
+        prob = x_numpy / x_numpy.sum(axis=-1, keepdims=True)
+        self.assertTrue(
+            np.allclose(
+                sample_prob, prob, rtol=0, atol=0.01),
+            "sample_prob: " + str(sample_prob) + "\nprob: " + str(prob))
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
index 8bb584be236..f98a5aff466 100755
--- a/python/paddle/tensor/__init__.py
+++ b/python/paddle/tensor/__init__.py
@@ -166,6 +166,7 @@ from .math import isfinite  #DEFINE_ALIAS
 from .math import isinf  #DEFINE_ALIAS
 from .math import isnan  #DEFINE_ALIAS
 from .math import prod  #DEFINE_ALIAS
+from .random import multinomial  #DEFINE_ALIAS
 from .random import standard_normal
 from .random import normal
 from .random import uniform  #DEFINE_ALIAS
diff --git a/python/paddle/tensor/random.py b/python/paddle/tensor/random.py
index b38a1d0f5b7..7740754fe99 100644
--- a/python/paddle/tensor/random.py
+++ b/python/paddle/tensor/random.py
@@ -25,6 +25,7 @@ from ..fluid.io import shuffle  #DEFINE_ALIAS
 
 __all__ = [
     'bernoulli',
+    'multinomial',
     'standard_normal',
     'normal',
     'uniform',
@@ -88,6 +89,47 @@ def bernoulli(x, name=None):
     return out
 
 
+def multinomial(x, num_samples=1, replacement=False, name=None):
+    """
+
+    
+    Examples:
+        .. code-block:: python
+
+        import paddle
+
+        paddle.disable_static()
+
+        x = paddle.rand([2, 3])
+        print(x.numpy())
+        # [[0.11272584 0.3890902  0.7730957 ]
+        # [0.10351662 0.8510418  0.63806665]]
+
+        out = paddle.bernoulli(x)
+        print(out.numpy())
+        # [[0. 0. 1.]
+        # [0. 0. 1.]]
+
+    """
+
+    if in_dygraph_mode():
+        return core.ops.multinomial(x, 'num_samples', num_samples,
+                                    'replacement', replacement)
+
+    check_variable_and_dtype(x, "x", ["float32", "float64"], "multinomial")
+
+    helper = LayerHelper("multinomial", **locals())
+    out = helper.create_variable_for_type_inference(
+        dtype=convert_np_dtype_to_dtype_('int64'))
+    helper.append_op(
+        type='multinomial',
+        inputs={"X": x},
+        outputs={'Out': out},
+        attrs={'num_samples': num_samples,
+               'replacement': replacement})
+    return out
+
+
 def gaussian(shape, mean=0.0, std=1.0, dtype=None, name=None):
     """
     This OP returns a Tensor filled with random values sampled from a Gaussian
-- 
GitLab