set_dtype and mod MultinomialFunctor

paddle/phi/infermeta/ternary.cc

@@ -86,6 +86,7 @@ void AddmmInferMeta(const MetaTensor& input,
 
   out->set_dims(make_ddim(output_dims));
   out->share_lod(input);
+  out->set_dtype(input.dtype());
 }
 
 }  // namespace phi

paddle/phi/infermeta/unary.cc

@@ -99,6 +99,7 @@ void CholeskyInferMeta(const MetaTensor& x, bool upper, MetaTensor* out) {
           dims[rank - 2],
           dims[rank - 1]));
   out->set_dims(x.dims());
+  out->set_dtype(x.dtype());
 }
 
 void CopyToInferMeta(const MetaTensor& x,
@@ -123,6 +124,7 @@ void IncrementInferMeta(const MetaTensor& x, float value, MetaTensor* out) {
                               product(x.dims())));
   out->set_dims(x.dims());
   out->share_lod(x);
+  out->set_dtype(x.dtype());
 }
 
 static phi::DDim ValidateShape(const std::vector<int64_t> shape,
@@ -297,6 +299,7 @@ void MultinomialInferMeta(const MetaTensor& x,
   out_dims[x_rank - 1] = num_samples;
 
   out->set_dims(make_ddim(out_dims));
+  out->set_dtype(DataType::INT64);
 }
 
 void ReshapeInferMeta(const MetaTensor& x,

paddle/phi/kernels/cpu/multinomial_kernel.cc

@@ -33,7 +33,8 @@ void MultinomialKernel(const Context& dev_ctx,
   const int64_t num_categories = in_dims[in_rank - 1];
   const int64_t num_distributions = in_rank > 1 ? in_dims[in_rank - 2] : 1;
 
-  funcs::MultinomialFunctor<T>(out_data,
+  funcs::MultinomialFunctor<T>(dev_ctx,
+                               out_data,
                                in_data,
                                num_samples,
                                replacement,

paddle/phi/kernels/funcs/multinomial_functor.h

@@ -14,14 +14,15 @@ limitations under the License. */
 
 #pragma once
 
-#include "paddle/fluid/framework/generator.h"
+#include "paddle/phi/core/device_context.h"
 #include "paddle/phi/core/enforce.h"
 
 namespace phi {
 namespace funcs {
 
-template <typename T>
-void MultinomialFunctor(int64_t* out_data,
+template <typename T, typename Context>
+void MultinomialFunctor(const Context& dev_ctx,
+                        int64_t* out_data,
                         const T* in_data,
                         const int64_t num_samples,
                         const bool replacement,
@@ -30,8 +31,7 @@ void MultinomialFunctor(int64_t* out_data,
   std::vector<T> cumulative_probs(num_categories);
 
   std::uniform_real_distribution<T> dist(0, 1);
-  auto gen_ptr = paddle::framework::DefaultCPUGenerator();
-  auto engine = gen_ptr->GetCPUEngine();
+  auto engine = dev_ctx.GetHostGenerator()->GetCPUEngine();
 
   for (int64_t i = 0; i < num_distributions; i++) {
     T probs_sum = 0;

paddle/phi/kernels/gpu/multinomial_kernel.cu

@@ -173,7 +173,8 @@ void MultinomialKernel(const Context& dev_ctx,
                cudaMemcpyDeviceToHost);
 #endif
 
-    funcs::MultinomialFunctor<T>(cpu_out_data,
+    funcs::MultinomialFunctor<T>(dev_ctx,
+                                 cpu_out_data,
                                  cpu_in_data,
                                  num_samples,
                                  replacement,