From 31c3312226f893968f82f991616f063cdad5eaca Mon Sep 17 00:00:00 2001
From: Zhou Wei <1183042833@qq.com>
Date: Wed, 27 Apr 2022 10:45:30 +0800
Subject: [PATCH] fix multinomial paddle_enforce bug (#42302)

---
 paddle/phi/kernels/gpu/multinomial_kernel.cu          | 11 +++++------
 .../fluid/tests/unittests/test_multinomial_op.py      |  8 ++++++++
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/paddle/phi/kernels/gpu/multinomial_kernel.cu b/paddle/phi/kernels/gpu/multinomial_kernel.cu
index ef6cd1323a9..21a506a840c 100644
--- a/paddle/phi/kernels/gpu/multinomial_kernel.cu
+++ b/paddle/phi/kernels/gpu/multinomial_kernel.cu
@@ -133,11 +133,10 @@ void MultinomialKernel(const Context& dev_ctx,
                        DenseTensor* out) {
   auto* in_data = x.data<T>();
   int64_t* out_data = dev_ctx.template Alloc<int64_t>(out);
-
   auto in_dims = x.dims();
-  int64_t in_rank = in_dims.size();
-  const int64_t num_categories = in_dims[in_rank - 1];
-  const int64_t num_distributions = in_rank > 1 ? in_dims[in_rank - 2] : 1;
+  int64_t dim_size = in_dims.size();
+  const int64_t num_categories = in_dims[dim_size - 1];
+  const int64_t num_distributions = dim_size > 1 ? in_dims[dim_size - 2] : 1;
 
   // If replacement is False, it's not a replaceable sample. Every category
   // can be used only once.
@@ -145,8 +144,8 @@ void MultinomialKernel(const Context& dev_ctx,
     int64_t in_data_numel = x.numel();
     int64_t out_data_numel = out->numel();
 
+    // Just use to PADDLE_ENFORCE error message
     T* cpu_in_data = new T[in_data_numel];
-    int64_t* cpu_out_data = new int64_t[out_data_numel];
 
 #ifdef PADDLE_WITH_HIP
     hipMemcpy(
@@ -160,7 +159,7 @@ void MultinomialKernel(const Context& dev_ctx,
     for (size_t i = 0; i < num_distributions; ++i) {
       int zero_num = 0;
       for (size_t j = 0; j < num_categories; ++j) {
-        T weight = cpu_in_data[i * num_distributions + j];
+        T weight = cpu_in_data[i * num_categories + j];
         PADDLE_ENFORCE_GE(
             weight,
             0,
diff --git a/python/paddle/fluid/tests/unittests/test_multinomial_op.py b/python/paddle/fluid/tests/unittests/test_multinomial_op.py
index ecde527523d..4dfc881d772 100644
--- a/python/paddle/fluid/tests/unittests/test_multinomial_op.py
+++ b/python/paddle/fluid/tests/unittests/test_multinomial_op.py
@@ -216,6 +216,14 @@ class TestMultinomialError(unittest.TestCase):
 
         self.assertRaises(ValueError, test_dim_less_than_1)
 
+        with self.assertRaises(ValueError):
+            y = paddle.multinomial(paddle.to_tensor([1., 2., -3.]))
+
+        with self.assertRaises(ValueError):
+            prob = paddle.rand([20, 1000])
+            prob[1:0] = 0
+            y = paddle.multinomial(prob)
+
 
 class TestRandomValue(unittest.TestCase):
     def test_fixed_random_number(self):
-- 
GitLab