fix multinomial paddle_enforce bug (#42302)

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,

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):