fix gamma bug

docs/api/api_python/ops/mindspore.ops.func_random_gamma.rst

@@ -1,7 +1,7 @@
-mindspore.ops.gamma
+mindspore.ops.random_gamma
 ====================
 
-.. py:function:: mindspore.ops.gamma(shape, alpha, seed=None)
+.. py:function:: mindspore.ops.random_gamma(shape, alpha, seed=None)
 
     根据伽马分布产生成随机数。
 

mindspore/ccsrc/plugin/device/cpu/kernel/eigen/gamma_cpu_kernel.cc

@@ -44,12 +44,25 @@ bool GammaCpuKernelMod::Init(const BaseOperatorPtr &base_operator, const std::ve
   alpha_shape_ = inputs[1]->GetShapeVector();
   alpha_dtype_ = inputs[1]->GetDtype();
   shape_dtype_ = inputs[0]->GetDtype();
+  shape_shape_ = inputs[0]->GetShapeVector();
 
   is_need_retrieve_output_shape_ = true;
 
   return true;
 }
 
+template <typename T>
+void GammaCpuKernelMod::InferShape(const std::vector<AddressPtr> &inputs) {
+  const auto *shape_value = reinterpret_cast<T *>(inputs[0]->addr);
+
+  for (int64_t i = 0; i < shape_shape_[0]; i++) {
+    output_shape_.emplace_back(static_cast<int64_t>(shape_value[i]));
+  }
+  for (size_t i = 0; i < alpha_shape_.size(); i++) {
+    output_shape_.emplace_back(alpha_shape_[i]);
+  }
+}
+
 int GammaCpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const std::vector<KernelTensorPtr> &inputs,
                               const std::vector<KernelTensorPtr> &outputs,
                               const std::map<uint32_t, tensor::TensorPtr> &inputsOnHost) {
@@ -62,9 +75,6 @@ int GammaCpuKernelMod::Resize(const BaseOperatorPtr &base_operator, const std::v
     return ret;
   }
 
-  shape_shape_ = inputs[0]->GetShapeVector();
-  alpha_shape_ = inputs[1]->GetShapeVector();
-
   return KRET_OK;
 }
 
@@ -74,7 +84,7 @@ void GammaCpuKernelMod::Generate(const std::vector<AddressPtr> &inputs, const st
   const auto *alpha_flat = reinterpret_cast<T *>(inputs[1]->addr);
   auto *samples_flat = reinterpret_cast<T *>(outputs[0]->addr);
 
-  int64_t num_samples = std::accumulate(output_shape_.begin(), output_shape_.end(), 0);
+  int64_t num_samples = std::accumulate(output_shape_.begin(), output_shape_.end(), 1, std::multiplies<int64_t>());
   if (num_samples == 0) {
     MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "' the sizes of output is zero.";
   }
@@ -92,11 +102,16 @@ void GammaCpuKernelMod::Generate(const std::vector<AddressPtr> &inputs, const st
 
   static constexpr int kReservedSamplesPerOutput = 256;
 
-  int64 num_alphas = std::accumulate(alpha_shape_.begin(), alpha_shape_.end(), 0);
+  int64 num_alphas = std::accumulate(alpha_shape_.begin(), alpha_shape_.end(), 1, std::multiplies<int64_t>());
+  if (num_alphas == 0) {
+    MS_LOG(EXCEPTION) << "For '" << kernel_name_ << "' the sizes of alpha is zero.";
+  }
+  int64_t sample_shape_per_al = num_samples / num_alphas;
 
-  PhiloxRandom rng = generator_.ReserveRandomOutputs(num_samples * num_alphas, kReservedSamplesPerOutput);
+  PhiloxRandom rng = generator_.ReserveRandomOutputs(num_samples, kReservedSamplesPerOutput);
 
-  auto DoWork = [num_samples, num_alphas, &rng, samples_flat, alpha_flat](int64 start_output, int64 limit_output) {
+  auto DoWork = [sample_shape_per_al, num_alphas, &rng, samples_flat, alpha_flat](int64 start_output,
+                                                                                  int64 limit_output) {
     using Eigen::numext::exp;
     using Eigen::numext::log;
     using Eigen::numext::pow;
@@ -107,7 +122,7 @@ void GammaCpuKernelMod::Generate(const std::vector<AddressPtr> &inputs, const st
     typename Uniform::ResultType uniform_res;
 
     for (int64 output_idx = start_output; output_idx < limit_output;) {
-      int64 alpha_idx = output_idx / num_samples;
+      int64 alpha_idx = output_idx / sample_shape_per_al;
       T *const samples_alpha_offset = samples_flat + alpha_idx;
       const double alpha_value = static_cast<double>(alpha_flat[alpha_idx]);
 
@@ -115,8 +130,8 @@ void GammaCpuKernelMod::Generate(const std::vector<AddressPtr> &inputs, const st
       if (alpha_value == static_cast<double>(1.0)) {
         //        ENABLE_FLOAT_EQUALITY_WARNING
         // Sample from an exponential distribution.
-        for (int64 sample_idx = output_idx % num_samples; sample_idx < num_samples && output_idx < limit_output;
-             sample_idx++, output_idx++) {
+        for (int64 sample_idx = output_idx % sample_shape_per_al;
+             sample_idx < sample_shape_per_al && output_idx < limit_output; sample_idx++, output_idx++) {
           PhiloxRandom gen = rng;
           gen.Skip(kReservedSamplesPerOutput * output_idx);
           int16 uniform_remaining = 0;
@@ -132,8 +147,8 @@ void GammaCpuKernelMod::Generate(const std::vector<AddressPtr> &inputs, const st
         const double cut = 1.0 / 3 / sqrt(su);
 
         // Compute the rest of the samples for the current alpha value.
-        for (int64 sample_idx = output_idx % num_samples; sample_idx < num_samples && output_idx < limit_output;
-             sample_idx++, output_idx++) {
+        for (int64 sample_idx = output_idx % sample_shape_per_al;
+             sample_idx < sample_shape_per_al && output_idx < limit_output; sample_idx++, output_idx++) {
           PhiloxRandom gen = rng;
           gen.Skip(kReservedSamplesPerOutput * output_idx);
           int16 norm_remaining = 0;
@@ -171,19 +186,7 @@ void GammaCpuKernelMod::Generate(const std::vector<AddressPtr> &inputs, const st
     }
   };
 #undef UNIFORM
-  ParallelLaunchAutoSearch(DoWork, num_alphas * num_samples, this, &parallel_search_info_);
-}
-
-template <typename T>
-void GammaCpuKernelMod::InferShape(const std::vector<AddressPtr> &inputs) {
-  const auto *shape_value = reinterpret_cast<T *>(inputs[0]->addr);
-
-  for (int64_t i = 0; i < shape_shape_[0]; i++) {
-    output_shape_.emplace_back(static_cast<int64_t>(shape_value[i]));
-  }
-  for (size_t i = 0; i < alpha_shape_.size(); i++) {
-    output_shape_.emplace_back(alpha_shape_[i]);
-  }
+  ParallelLaunchAutoSearch(DoWork, num_alphas * sample_shape_per_al, this, &parallel_search_info_);
 }
 
 bool GammaCpuKernelMod::Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,

mindspore/python/mindspore/ops/operations/random_ops.py

@@ -321,8 +321,8 @@ class Gamma(PrimitiveWithInfer):
         >>> shape = (3, 1, 2)
         >>> alpha = Tensor(np.array([[3, 4], [5, 6]]), mstype.float32)
         >>> beta = Tensor(np.array([1.0]), mstype.float32)
-        >>> random_gamma = ops.RandomGamma(seed=3)
-        >>> output = random_gamma(shape, alpha, beta)
+        >>> gamma = ops.Gamma(seed=3)
+        >>> output = gamma(shape, alpha, beta)
         >>> result = output.shape
         >>> print(result)
         (3, 2, 2)

tests/st/ops/cpu/test_random_gamma_op.py

@@ -71,10 +71,10 @@ def test_random_gamma_functional(dtype):
     """
 
     ms.set_context(mode=ms.GRAPH_MODE, device_target='CPU')
-    shape = Tensor(np.array([3, 2]), ms.int32)
+    shape = Tensor(np.array([10, 10]), ms.int32)
     alpha = Tensor(np.array([[3, 4], [5, 6]]), dtype)
     output = F.random_gamma(shape=shape, alpha=alpha, seed=2)
-    expect = np.array([3, 2, 2, 2])
+    expect = np.array([10, 10, 2, 2])
 
     print(output)
     assert (output.shape == expect).all()