!17135 complete forward and backward of SparseToDense and check of SparseTensorDenseMatmul

From: @zhangbuxue
Reviewed-by: 
Signed-off-by:

mindspore/ccsrc/backend/kernel_compiler/cpu/gathernd_cpu_kernel.cc

@@ -19,7 +19,8 @@
 
 namespace mindspore {
 namespace kernel {
-void GatherNdCPUKernel::InitKernel(const CNodePtr &kernel_node) {
+template <typename T>
+void GatherNdCPUKernel<T>::InitKernel(const CNodePtr &kernel_node) {
   input_shapes_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
   indices_shapes_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
   output_shapes_ = AnfAlgo::GetOutputInferShape(kernel_node, 0);
@@ -56,28 +57,14 @@ void GatherNdCPUKernel::InitKernel(const CNodePtr &kernel_node) {
   }
 }
 
-bool GatherNdCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &,
-                               const std::vector<kernel::AddressPtr> &outputs) {
-  if (dtype_ == kNumberTypeInt32) {
-    return LaunchKernel<int32_t>(inputs, outputs);
-  } else if (dtype_ == kNumberTypeInt64) {
-    return LaunchKernel<int64_t>(inputs, outputs);
-  } else if (dtype_ == kNumberTypeFloat32) {
-    return LaunchKernel<float>(inputs, outputs);
-  } else if (dtype_ == kNumberTypeFloat64) {
-    return LaunchKernel<double>(inputs, outputs);
-  } else {
-    MS_LOG(EXCEPTION) << "Only support int, float, but actual data type is " << TypeIdLabel(dtype_);
-  }
-}
-
 template <typename T>
-bool GatherNdCPUKernel::LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs) {
+bool GatherNdCPUKernel<T>::Launch(const std::vector<kernel::AddressPtr> &inputs,
+                                  const std::vector<kernel::AddressPtr> &,
+                                  const std::vector<kernel::AddressPtr> &outputs) {
   auto input_addr = reinterpret_cast<T *>(inputs[0]->addr);
   auto indices_addr = reinterpret_cast<int *>(inputs[1]->addr);
   auto output_addr = reinterpret_cast<T *>(outputs[0]->addr);
 
-  //
   size_t output_dim0 = dims_[0];
   size_t output_dim1 = dims_[1];
   size_t indices_dim1 = dims_[2];

mindspore/ccsrc/backend/kernel_compiler/cpu/gathernd_cpu_kernel.h

@@ -22,6 +22,7 @@
 
 namespace mindspore {
 namespace kernel {
+template <typename T>
 class GatherNdCPUKernel : public CPUKernel {
  public:
   GatherNdCPUKernel() = default;
@@ -32,9 +33,6 @@ class GatherNdCPUKernel : public CPUKernel {
   bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
               const std::vector<AddressPtr> &outputs) override;
 
-  template <typename T>
-  bool LaunchKernel(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &outputs);
-
  private:
   std::vector<size_t> input_shapes_;
   std::vector<size_t> indices_shapes_;
@@ -47,7 +45,18 @@ class GatherNdCPUKernel : public CPUKernel {
   TypeId dtype_{kTypeUnknown};
 };
 
-MS_REG_CPU_KERNEL(GatherNd, KernelAttr(), GatherNdCPUKernel);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, bool);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, int8_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, int16_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, int32_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, int64_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, uint8_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, uint16_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, uint32_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, uint64_t);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, float);
+MS_REG_CPU_KERNEL_T(GatherNd, KernelAttr(), GatherNdCPUKernel, double);
+
 }  // namespace kernel
 }  // namespace mindspore
 

mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.cc

@@ -24,13 +24,21 @@ template <typename I, typename T>
 void SparseTensorDenseMatmulCPUKernel<I, T>::InitKernel(const CNodePtr &kernel_node) {
   adj_st_ = AnfAlgo::GetNodeAttr<bool>(kernel_node, ADJ_ST);
   adj_dt_ = AnfAlgo::GetNodeAttr<bool>(kernel_node, ADJ_dT);
-  output_shape_ = AnfAlgo::GetOutputInferShape(kernel_node, 0);
+  auto indices_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
-  output_size_ = std::accumulate(output_shape_.begin(), output_shape_.end(), size_t(1), std::multiplies<size_t>());
+  if (indices_shape.size() != 2 && indices_shape[1] != 2) {
-  auto values_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
+    MS_LOG(EXCEPTION)
-  if (values_shape.size() != 1) {
+      << "SparseTensorDenseMatmul requires 'indices' should be a 2-D Tensor and the second dimension length "
-    MS_LOG(EXCEPTION) << "SparseTensorDenseMatmul requires the values must be a 1-D tensor, but got "
+         "should be 2, but got 'indices' shape: "
-                      << values_shape.size() << "-D";
+      << indices_shape;
   }
+  auto values_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
+  if (values_shape.size() != 1 || values_shape[0] != indices_shape[0]) {
+    MS_LOG(EXCEPTION)
+      << "SparseTensorDenseMatmul requires 'value's should be a 1-D Tensor and the first dimension length should be "
+         "equal to the first dimension length of 'indices', but got 'values' shape: "
+      << values_shape;
+  }
+  output_shape_ = AnfAlgo::GetOutputInferShape(kernel_node, 0);
   values_size_ = values_shape[0];
   b_shape_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 3);
 }
@@ -43,21 +51,20 @@ bool SparseTensorDenseMatmulCPUKernel<I, T>::Launch(const std::vector<kernel::Ad
   auto a_values = reinterpret_cast<T *>(inputs[1]->addr);
   auto b = reinterpret_cast<T *>(inputs[3]->addr);
   auto out = reinterpret_cast<T *>(outputs[0]->addr);
-  memset(out, 0, output_size_);
+  const size_t output_length = outputs[0]->size / sizeof(T);
+  memset(out, 0, output_length);
 
   const size_t out_dim_0 = output_shape_[0];
   const size_t out_dim_1 = output_shape_[1];
   const size_t b_dim_0 = b_shape_[0];
   const size_t b_dim_1 = b_shape_[1];
   const size_t same_dim = adj_dt_ ? b_dim_1 : b_dim_0;
-
   for (size_t i = 0; i < values_size_; ++i) {
     const int row = adj_st_ ? a_indices[i * 2 + 1] : a_indices[i * 2];
     const int col = adj_st_ ? a_indices[i * 2] : a_indices[i * 2 + 1];
     if (row >= SizeToInt(out_dim_0) || row < 0 || col >= SizeToInt(same_dim) || col < 0) {
-      MS_LOG(ERROR) << "The indices including out of bounds index, row range: [0, " << out_dim_0 << "), col range: [0, "
+      MS_EXCEPTION(ValueError) << "The indices including out of bounds index, row range: [0, " << out_dim_0
-                    << same_dim << "), but got row: " << row << ", col: " << col;
+                               << "), col range: [0, " << same_dim << "), but got row: " << row << ", col: " << col;
-      return false;
     }
 
     for (size_t n = 0; n < out_dim_1; ++n) {

mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_tensor_dense_matmul_cpu_kernel.h

@@ -42,68 +42,6 @@ class SparseTensorDenseMatmulCPUKernel : public CPUKernel {
   bool adj_st_{false};
   bool adj_dt_{false};
 };
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeBool)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeBool)
-                        .AddOutputAttr(kNumberTypeBool),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, bool);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt8)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt8)
-                        .AddOutputAttr(kNumberTypeUInt8),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, uint8_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt16)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt16)
-                        .AddOutputAttr(kNumberTypeUInt16),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, uint16_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt32)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt32)
-                        .AddOutputAttr(kNumberTypeUInt32),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, uint32_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt64)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt64)
-                        .AddOutputAttr(kNumberTypeUInt64),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, uint64_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt8)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt8)
-                        .AddOutputAttr(kNumberTypeInt8),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, int8_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt16)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt16)
-                        .AddOutputAttr(kNumberTypeInt16),
-                      SparseTensorDenseMatmulCPUKernel, int32_t, int16_t);
 
 MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
                       KernelAttr()
@@ -141,105 +79,6 @@ MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
                         .AddOutputAttr(kNumberTypeFloat64),
                       SparseTensorDenseMatmulCPUKernel, int32_t, double);
 
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeBool)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeBool)
-                        .AddOutputAttr(kNumberTypeBool),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, bool);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeUInt8)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt8)
-                        .AddOutputAttr(kNumberTypeUInt8),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, uint8_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeUInt16)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt16)
-                        .AddOutputAttr(kNumberTypeUInt16),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, uint16_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeUInt32)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt32)
-                        .AddOutputAttr(kNumberTypeUInt32),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, uint32_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeUInt64)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt64)
-                        .AddOutputAttr(kNumberTypeUInt64),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, uint64_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt8)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt8)
-                        .AddOutputAttr(kNumberTypeInt8),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, int8_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt16)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt16)
-                        .AddOutputAttr(kNumberTypeInt16),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, int16_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeInt32),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, int32_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeInt64),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, int64_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeFloat32)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeFloat32)
-                        .AddOutputAttr(kNumberTypeFloat32),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, float);
-
-MS_REG_CPU_KERNEL_T_S(SparseTensorDenseMatmul,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeFloat64)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeFloat64)
-                        .AddOutputAttr(kNumberTypeFloat64),
-                      SparseTensorDenseMatmulCPUKernel, int64_t, double);
-
 }  // namespace kernel
 }  // namespace mindspore
 #endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_RMSPROP_CPU_KERNEL_H_

mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_to_dense_cpu_kernal.cc

@@ -23,24 +23,20 @@ namespace kernel {
 template <typename I, typename T>
 void SparseToDenseCPUKernel<I, T>::InitKernel(const CNodePtr &kernel_node) {
   CheckParam(kernel_node);
-  indices_shape_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
+  auto indices_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 0);
-  values_shape_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
+  if (indices_shape.size() != 2) {
-  dense_shape_shape_ = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 2);
+    MS_LOG(EXCEPTION) << "SparseToDense requires 'indices' should be a 2-D Tensor, but got " << indices_shape.size()
+                      << "-D";
+  }
+  auto values_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, 1);
+  if (values_shape.size() != 1 || values_shape[0] != indices_shape[0]) {
+    MS_LOG(EXCEPTION)
+      << "SparseToDense requires 'values' should be a 1-D Tensor and the first dimension length should be "
+         "equal to the 'indices' first dimension length, but got 'values' shape: "
+      << values_shape;
+  }
+  values_size_ = values_shape[0];
   output_shape_ = AnfAlgo::GetOutputInferShape(kernel_node, 0);
-  if (!indices_shape_.size() || !values_shape_.size() || !output_shape_.size()) {
-    MS_LOG(EXCEPTION) << "Input NULL";
-  }
-  if (indices_shape_.size() > 2 || indices_shape_[0] != values_shape_[0]) {
-    MS_LOG(EXCEPTION) << "Input Error";
-  }
-}
-
-size_t DenseGetTensorLen(const std::vector<size_t> &shape) {
-  size_t len = 1;
-  for (size_t i = 0; i < shape.size(); i++) {
-    len *= shape[i];
-  }
-  return len;
 }
 
 template <typename I, typename T>
@@ -50,36 +46,26 @@ bool SparseToDenseCPUKernel<I, T>::Launch(const std::vector<kernel::AddressPtr>
   auto indices_addr = reinterpret_cast<I *>(inputs[0]->addr);
   auto values_addr = reinterpret_cast<T *>(inputs[1]->addr);
   auto output_addr = reinterpret_cast<T *>(outputs[0]->addr);
+  const size_t output_length = outputs[0]->size / sizeof(T);
+  memset(output_addr, 0, output_length);
 
-  size_t output_len = DenseGetTensorLen(output_shape_);
+  size_t rank = output_shape_.size();
-  memset(output_addr, 0, output_len * sizeof(T));
+  for (size_t i = 0; i < values_size_; ++i) {
-  std::vector<size_t> cargo(output_shape_.size(), 0);
-
-  size_t i = output_shape_.size() - 1;
-  switch (indices_shape_.size()) {
-    case 1:
-      for (i = 0; i < indices_shape_[0]; i++) {
-        output_addr[indices_addr[i]] = values_addr[i];
-      }
-      break;
-
-    case 2:
-      cargo[i] = 1;
-      for (; i >= 1; i--) {
-        cargo[i - 1] = cargo[i] * output_shape_[i];
-      }
-      for (i = 0; i < indices_shape_[0]; i++) {
     size_t out_index = 0;
-        for (size_t j = 0; j < indices_shape_[1]; j++) {
+    for (size_t j = 0; j < rank; j++) {
-          out_index += (*(indices_addr + i * indices_shape_[1] + j)) * cargo[j];
+      int index = indices_addr[i * rank + j];
+      if (index >= SizeToInt(output_shape_[j]) || index < 0) {
+        MS_EXCEPTION(ValueError) << "The " << i << "th value in " << j << "th dimension index: " << index
+                                 << " out of bounds: [0, " << output_shape_[j] << ")";
+      }
+      size_t count = 1;
+      for (size_t k = j + 1; k < rank; k++) {
+        count *= output_shape_[k];
+      }
+      out_index += SizeToInt(index) * count;
     }
     output_addr[out_index] = values_addr[i];
   }
-      break;
-
-    default:
-      break;
-  }
   return true;
 }
 
@@ -87,11 +73,11 @@ template <typename I, typename T>
 void SparseToDenseCPUKernel<I, T>::CheckParam(const CNodePtr &kernel_node) {
   size_t input_num = AnfAlgo::GetInputTensorNum(kernel_node);
   if (input_num != 3) {
-    MS_LOG(EXCEPTION) << "Input number is " << input_num << ", but SparseToDenseCPUKernel needs 3 input.";
+    MS_LOG(EXCEPTION) << "SparseToDense needs 3 inputs, but got " << input_num;
   }
   size_t output_num = AnfAlgo::GetOutputTensorNum(kernel_node);
   if (output_num != 1) {
-    MS_LOG(EXCEPTION) << "Output number is " << output_num << ", but SparseToDenseCPUKernel needs 1 output.";
+    MS_LOG(EXCEPTION) << "SparseToDense should have 2 outputs, but got " << output_num;
   }
 }
 }  // namespace kernel

mindspore/ccsrc/backend/kernel_compiler/cpu/sparse_to_dense_cpu_kernal.h

@@ -37,19 +37,17 @@ class SparseToDenseCPUKernel : public CPUKernel {
 
  private:
   void CheckParam(const CNodePtr &kernel_node);
-  std::vector<size_t> indices_shape_;
-  std::vector<size_t> values_shape_;
-  std::vector<size_t> dense_shape_shape_;
   std::vector<size_t> output_shape_;
+  size_t values_size_{0};
 };
 
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
                         .AddInputAttr(kNumberTypeInt32)
+                        .AddInputAttr(kNumberTypeBool)
                         .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt32)
+                        .AddOutputAttr(kNumberTypeBool),
-                        .AddOutputAttr(kNumberTypeInt32),
+                      SparseToDenseCPUKernel, int32_t, bool);
-                      SparseToDenseCPUKernel, int32_t, int32_t);
 
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
@@ -59,14 +57,6 @@ MS_REG_CPU_KERNEL_T_S(SparseToDense,
                         .AddOutputAttr(kNumberTypeInt8),
                       SparseToDenseCPUKernel, int32_t, int8_t);
 
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt8)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeUInt8),
-                      SparseToDenseCPUKernel, int32_t, uint8_t);
-
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
                         .AddInputAttr(kNumberTypeInt32)
@@ -78,10 +68,10 @@ MS_REG_CPU_KERNEL_T_S(SparseToDense,
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
                         .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeUInt16)
                         .AddInputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeUInt16),
+                        .AddInputAttr(kNumberTypeInt32)
-                      SparseToDenseCPUKernel, int32_t, uint16_t);
+                        .AddOutputAttr(kNumberTypeInt32),
+                      SparseToDenseCPUKernel, int32_t, int32_t);
 
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
@@ -91,6 +81,38 @@ MS_REG_CPU_KERNEL_T_S(SparseToDense,
                         .AddOutputAttr(kNumberTypeInt64),
                       SparseToDenseCPUKernel, int32_t, int64_t);
 
+MS_REG_CPU_KERNEL_T_S(SparseToDense,
+                      KernelAttr()
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddInputAttr(kNumberTypeUInt8)
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddOutputAttr(kNumberTypeUInt8),
+                      SparseToDenseCPUKernel, int32_t, uint8_t);
+
+MS_REG_CPU_KERNEL_T_S(SparseToDense,
+                      KernelAttr()
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddInputAttr(kNumberTypeUInt16)
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddOutputAttr(kNumberTypeUInt16),
+                      SparseToDenseCPUKernel, int32_t, uint16_t);
+
+MS_REG_CPU_KERNEL_T_S(SparseToDense,
+                      KernelAttr()
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddInputAttr(kNumberTypeUInt32)
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddOutputAttr(kNumberTypeUInt32),
+                      SparseToDenseCPUKernel, int32_t, uint32_t);
+
+MS_REG_CPU_KERNEL_T_S(SparseToDense,
+                      KernelAttr()
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddInputAttr(kNumberTypeUInt64)
+                        .AddInputAttr(kNumberTypeInt32)
+                        .AddOutputAttr(kNumberTypeUInt64),
+                      SparseToDenseCPUKernel, int32_t, uint64_t);
+
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
                         .AddInputAttr(kNumberTypeInt32)
@@ -102,9 +124,9 @@ MS_REG_CPU_KERNEL_T_S(SparseToDense,
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
                       KernelAttr()
                         .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeFloat)
+                        .AddInputAttr(kNumberTypeFloat32)
                         .AddInputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeFloat),
+                        .AddOutputAttr(kNumberTypeFloat32),
                       SparseToDenseCPUKernel, int32_t, float);
 
 MS_REG_CPU_KERNEL_T_S(SparseToDense,
@@ -115,93 +137,6 @@ MS_REG_CPU_KERNEL_T_S(SparseToDense,
                         .AddOutputAttr(kNumberTypeFloat64),
                       SparseToDenseCPUKernel, int32_t, double);
 
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeBool)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddOutputAttr(kNumberTypeBool),
-                      SparseToDenseCPUKernel, int32_t, bool);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt32)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeInt32),
-                      SparseToDenseCPUKernel, int64_t, int32_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeInt64),
-                      SparseToDenseCPUKernel, int64_t, int64_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt8)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeInt8),
-                      SparseToDenseCPUKernel, int64_t, int8_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeUInt8)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeUInt8),
-                      SparseToDenseCPUKernel, int64_t, uint8_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeInt16)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeInt16),
-                      SparseToDenseCPUKernel, int64_t, int16_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeUInt16)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeUInt16),
-                      SparseToDenseCPUKernel, int64_t, uint16_t);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeFloat16)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeFloat16),
-                      SparseToDenseCPUKernel, int64_t, float16);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeFloat)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeFloat),
-                      SparseToDenseCPUKernel, int64_t, float);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeFloat64)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeFloat64),
-                      SparseToDenseCPUKernel, int64_t, double);
-
-MS_REG_CPU_KERNEL_T_S(SparseToDense,
-                      KernelAttr()
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddInputAttr(kNumberTypeBool)
-                        .AddInputAttr(kNumberTypeInt64)
-                        .AddOutputAttr(kNumberTypeBool),
-                      SparseToDenseCPUKernel, int64_t, bool);
 }  // namespace kernel
 }  // namespace mindspore
 #endif  // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_CPU_SPARSETODENSE_CPU_KERNEL_H_

mindspore/nn/sparse/sparse.py

@@ -19,19 +19,19 @@ from ..cell import Cell
 
 class SparseToDense(Cell):
     """
-    Converts a sparse representation into a dense tensor.
+    Converts a sparse tensor into dense.
 
     Not yet supported by any backend at the moment.
 
     Inputs:
-        sparse_tensor (SparseTensor): the sparse tensor to convert.
+        - **sparse_tensor** (SparseTensor): the sparse tensor to convert.
 
     Outputs:
-        Tensor, the tensor converted.
+        Tensor, converted from sparse tensor.
 
     Raises:
-        TypeError: If `sparse_tensor.indices` is neither int32 nor int64.
+        TypeError: If the`sparse_tensor.indices` data type is neither int32 nor int64.
-        TypeError: If 'sparse_tensor.values' is not a Number.
+        TypeError: If the 'sparse_tensor.values' data type is not a Number or bool.
         TypeError: If 'sparse_tensor.dense_shape' is not a tuple.
 
     Supported Platforms:
@@ -52,6 +52,7 @@ class SparseToDense(Cell):
          [0 0 2 0]
          [0 0 0 0]]
     """
+
     def __init__(self):
         super(SparseToDense, self).__init__()
         self.sparse_to_dense = P.SparseToDense()
@@ -61,59 +62,65 @@ class SparseToDense(Cell):
                                     sparse_tensor.values,
                                     sparse_tensor.dense_shape)
 
+
 class SparseTensorDenseMatmul(Cell):
     """
-    Multiply SparseTensor(of rank 2) "A" by dense tensor.
+    Multiplies sparse matrix `a` and dense matrix `b`.
-    The shape of sparse tensor is :math:`(N, C)`, and the shape of dense tensor is :math:`(C, M)`, then the shape of
+    The rank of sparse matrix and dense matrix must equal to `2`.
-    output tensor is :math:`(N, M)`.The output data type is the same as "values".
 
     Args:
-        - *adjoint_st** (Bool) - If true, SparseTensor is transposed before multiplication. Default: False.
+        - *adjoint_st** (bool) - If true, sparse tensor is transposed before multiplication. Default: False.
-        - *adjoint_dt** (Bool) - If true, DenseTensor is transposed before multiplication. Default: False.
+        - *adjoint_dt** (bool) - If true, dense tensor is transposed before multiplication. Default: False.
 
     Inputs:
-        - **indices** (Tensor) - A 2D tensor with shape [N, rank], containing the indices of the nonzero values. The
+        - **indices** (Tensor) - A 2-D Tensor, represents the position of the element in the sparse tensor.
-          indices of sparse representation, support int32/int64.
+            Support int32, int64, each element value should be non-negative. The shape is :math:`(n, 2)`.
-        - **values** (Tensor) - A 1D tensor with shape [N] containing all nonzero values. Values corresponding to
+        - **values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`.
-          each row of indices.
+            Support float16, float32, float64, int32, int64. The shape should be :math:`(n,).
-        - **dense_shape** (tuple) - A 1D tuple with shape (N, M), specifying the shape of the tensor. An int tuple
+        - **sparse_shape** (tuple) - A positive int tuple which specifies the shape of sparse tensor,
-          which specifies the shape of dense tensor. The dense_shape is : math:`(N, C)`. If `adjoint_st` is True,
+            should have 2 elements, represent sparse tensor shape is :math:`(N, C)`.
-          its shape must be :math:`(N, C)` after transpose.
+        - **dense** (Tensor) - A 2-D Tensor, the dtype is same as `values`.
-        - **dense** (Tensor) - Dense Matrix. The shape of the tensor is :math:`(C, M)`. If
+            If `adjoint_st` is False and `adjoint_dt` is False, the shape must be :math:`(C, M)`.
-          `adjoint_dt` is True, its shape must be :math:`(C, M)` after transpose.
+            If `adjoint_st` is False and `adjoint_dt` is True, the shape must be :math:`(M, C)`.
+            If `adjoint_st` is True and `adjoint_dt` is False, the shape must be :math:`(N, M)`.
+            If `adjoint_st` is True and `adjoint_dt` is True, the shape must be :math:`(M, N)`.
 
-    Returns:
+    Outputs:
-        Tensor, the shape of tensor  is :math:`(N, M)`.The output data type is the same as "values".
+        Tensor, the dtype is the same as `values`.
+        If `adjoint_st` is False, the shape is :math:`(N, M)`.
+        If `adjoint_st` is True, the shape is :math:`(C, M)`.
+
+    Raises:
+        TypeError: If the type of `adjoint_st` or `adjoint_dt` is not bool, or the dtype of `indices`,
+            dtype of `values` and dtype of `dense` don't meet the parameter description.
+        ValueError: If `sparse_shape`, shape of `indices, shape of `values`,
+            and shape of `dense` don't meet the parameter description.
 
     Supported Platforms:
         ``CPU``
 
     Examples:
-        >>> class NetSparseDenseMatmul(nn.Cell):
+        >>> import mindspore as ms
-        ...     def __init__(self):
+        >>> from mindspore import Tensor
-        ...         super(NetSparseDenseMatmul, self).__init__()
+        >>> from mindspore import nn
-        ...         self.matmul = nn.SparseTensorDenseMatmul()
-        ...
-        ...     def construct(self, indices, values, dens_shape, dt):
-        ...         return self.matmul(indices, values, dens_shape, dt)
-        ...
         >>> indices = Tensor([[0, 1], [1, 2]], dtype=ms.int32)
         >>> values = Tensor([1, 2], dtype=ms.float32)
-        >>> dense_shape = (3, 4)
+        >>> sparse_shape = (3, 4)
-        >>> dsMatrix = Tensor([[1, 1], [2, 2], [3, 3], [4, 4]], dtype=ms.float32)
+        >>> dense = Tensor([[1, 1], [2, 2], [3, 3], [4, 4]], dtype=ms.float32)
-        >>> test_SparseDenseMatmul = NetSparseDenseMatmul()
+        >>> sparse_dense_matmul = nn.SparseTensorDenseMatmul()
-        >>> out = test_SparseDenseMatmul(indices, values, dense_shape, dsMatrix)
+        >>> out = sparse_dense_matmul(indices, values, sparse_shape, dense)
         >>> print(out)
         [[2 2]
          [0 6]
          [6 0]]
     """
+
     def __init__(self, adjoint_st=False, adjoint_dt=False):
         """Initialize SparseTensorDenseMatmul"""
         super(SparseTensorDenseMatmul, self).__init__()
-        self.adjst = adjoint_st
+        self.adj_st = adjoint_st
-        self.adjdt = adjoint_dt
+        self.adj_dt = adjoint_dt
-        self.matmul = P.SparseTensorDenseMatmul(adjoint_st=self.adjst, adjoint_dt=self.adjdt)
+        self.sparse_dense_matmul = P.SparseTensorDenseMatmul(adjoint_st=self.adj_st, adjoint_dt=self.adj_dt)
 
-    def construct(self, indices, values, dense_shape, dense):
+    def construct(self, indices, values, sparse_shape, dense):
-        return self.matmul(indices, values, dense_shape, dense)
+        return self.sparse_dense_matmul(indices, values, sparse_shape, dense)

mindspore/ops/_grad/grad_sparse.py

@@ -54,7 +54,7 @@ def get_bprop_sparse_to_dense(self):
     """Generate bprop for SparseToDense"""
 
     def bprop(indices, values, dense_shape, out, dout):
-        return zeros_like(indices), dout, zeros_like(dense_shape)
+        return zeros_like(indices), F.gather_nd(dout, indices), zeros_like(dense_shape)
 
     return bprop
 

mindspore/ops/_op_impl/cpu/gather_nd.py

@@ -20,8 +20,15 @@ gather_nd_op_info = CpuRegOp("GatherNd") \
     .input(0, "x1", "required") \
     .input(1, "x2", "required") \
     .output(0, "y", "required") \
+    .dtype_format(DataType.BOOL_Default, DataType.I32_Default, DataType.BOOL_Default)  \
+    .dtype_format(DataType.I8_Default, DataType.I32_Default, DataType.I8_Default)      \
+    .dtype_format(DataType.I16_Default, DataType.I32_Default, DataType.I16_Default)    \
     .dtype_format(DataType.I32_Default, DataType.I32_Default, DataType.I32_Default)    \
     .dtype_format(DataType.I64_Default, DataType.I32_Default, DataType.I64_Default)    \
+    .dtype_format(DataType.U8_Default, DataType.I32_Default, DataType.U8_Default)      \
+    .dtype_format(DataType.U16_Default, DataType.I32_Default, DataType.U16_Default)    \
+    .dtype_format(DataType.U32_Default, DataType.I32_Default, DataType.U32_Default)    \
+    .dtype_format(DataType.U64_Default, DataType.I32_Default, DataType.U64_Default)    \
     .dtype_format(DataType.F32_Default, DataType.I32_Default, DataType.F32_Default)    \
     .dtype_format(DataType.F64_Default, DataType.I32_Default, DataType.F64_Default)    \
     .get_op_info()

mindspore/ops/operations/sparse_ops.py

@@ -28,12 +28,19 @@ class SparseToDense(PrimitiveWithInfer):
     Converts a sparse representation into a dense tensor.
 
     Inputs:
-        - **indices** (Tensor) - The indices of sparse representation.
+        - **indices** (Tensor) - A 2-D Tensor, represents the position of the element in the sparse tensor.
-        - **values** (Tensor) - Values corresponding to each row of indices.
+            Support int32, int64, each element value should be a non-negative int number. The shape is :math:`(n, 2)`.
-        - **dense_shape** (tuple) - An int tuple which specifies the shape of dense tensor.
+        - **values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`.
+            The shape should be :math:`(n,).
+        - **sparse_shape** (tuple(int)) - A positive int tuple which specifies the shape of sparse tensor,
+            should have 2 elements, represent sparse tensor shape is :math:`(N, C)`.
 
     Returns:
-        Tensor, the shape of tensor is `dense_shape`.
+        Tensor, converted from sparse tensor. The dtype is same as `values`, and the shape is `sparse_shape`.
+
+    Raises:
+        TypeError: If the dtype of `indices` is neither int32 nor int64.
+        ValueError: If `sparse_shape`, shape of `indices and shape of `values` don't meet the parameter description.
 
     Supported Platforms:
         ``CPU``
@@ -41,8 +48,13 @@ class SparseToDense(PrimitiveWithInfer):
     Examples:
         >>> indices = Tensor([[0, 1], [1, 2]])
         >>> values = Tensor([1, 2], dtype=ms.float32)
-        >>> dense_shape = (3, 4)
+        >>> sparse_shape = (3, 4)
-        >>> out = ops.SparseToDense()(indices, values, dense_shape)
+        >>> sparse_to_dense = ops.SparseToDense()
+        >>> out = sparse_to_dense(indices, values, sparse_shape)
+        >>> print(out)
+        [[0 1 0 0]
+         [0 0 2 0]
+         [0 0 0 0]]
     """
 
     @prim_attr_register
@@ -50,10 +62,28 @@ class SparseToDense(PrimitiveWithInfer):
         """Initialize index_select"""
         self.init_prim_io_names(inputs=['indices', 'values', 'dense_shape'], outputs=['output'])
 
-    def __infer__(self, indices, values, dense_shape):
+    def __infer__(self, indices, values, sparse_shape):
-        validator.check_subclass("indices", indices['dtype'], mstype.tensor, self.name)
+        validator.check_tensor_dtype_valid('indices', indices['dtype'], [mstype.int32, mstype.int64], self.name)
-        validator.check_subclass("values", values['dtype'], mstype.tensor, self.name)
+        validator.check_tensor_dtype_valid('values', values['dtype'], mstype.number_type + (mstype.bool_,), self.name)
-        out = {'shape': dense_shape['value'],
+        indices_shape = indices['shape']
+        if len(indices_shape) != 2:
+            raise ValueError("SparseToDense requires 'indices' must be a 2-D Tensor, "
+                             f"but got 'indices' shape: {indices_shape}")
+        values_shape = values['shape']
+        if len(values_shape) != 1 or values_shape[0] != indices_shape[0]:
+            raise ValueError("SparseToDense requires 'values' must be a 1-D Tensor and "
+                             "the first dimension length must be equal to the first dimension length of 'indices', "
+                             f"but got 'indices' shape: {indices_shape}, 'values' shape: {values_shape}")
+        sparse_shape_v = sparse_shape['value']
+        for i in sparse_shape_v:
+            if isinstance(i, bool) or not isinstance(i, int) or i <= 0:
+                raise ValueError("SparseToDense requires all elements in 'sparse_shape' must be "
+                                 f"positive int number, but got 'sparse_shape': {sparse_shape_v}")
+        if len(sparse_shape_v) != indices_shape[1]:
+            raise ValueError("SparseToDense requires the 'sparse_shape' length should be equal to the 'indices' "
+                             "second dimension length, but got the 'indices' second dimension length: "
+                             f"{indices_shape[1]}, 'sparse_shape' length: {len(sparse_shape_v)}")
+        out = {'shape': sparse_shape['value'],
                'dtype': values['dtype'],
                'value': None}
         return out
@@ -61,31 +91,36 @@ class SparseToDense(PrimitiveWithInfer):
 
 class SparseTensorDenseMatmul(PrimitiveWithInfer):
     """
-    Multiply SparseTensor(of rank 2) "A" by dense tensor.
+    Multiplies sparse matrix `A` by dense matrix `B`.
-    The shape of sparse tensor is :math:`(N, C)`, and the shape of dense tensor is :math:`(C, M)`, then the shape of
+    The rank of sparse matrix and dense matrix must equal to `2`.
-    output tensor is :math:`(N, M)`.The output data type is the same as "values".
-    tensors.
 
     Args:
-        - *adjoint_st** (Bool) - If true, SparseTensor is transposed before multiplication. Default: False.
+        - *adjoint_st** (bool) - If true, sparse tensor is transposed before multiplication. Default: False.
-        - *adjoint_dt** (Bool) - If true, DenseTensor is transposed before multiplication. Default: False.
+        - *adjoint_dt** (bool) - If true, dense tensor is transposed before multiplication. Default: False.
 
     Inputs:
-        - **indices** (Tensor) - The indices of sparse representation, support int32/int64.
+        - **indices** (Tensor) - A 2-D Tensor, represents the position of the element in the sparse tensor.
-        - **values** (Tensor) - Values corresponding to each row of indices.
+            Support int32, int64, each element value should be a non-negative int number. The shape is :math:`(n, 2)`.
-        - **dense_shape** (tuple) - An int tuple which specifies the shape of dense tensor. The dense_shape is :
+        - **values** (Tensor) - A 1-D Tensor, represents the value corresponding to the position in the `indices`.
-          math:`(N, C)`. If `adjoint_st` is True, its shape must be :math:`(N, C)` after transpose.
+            Support float16, float32, float64, int32, int64. The shape should be :math:`(n,).
-        - **dense** (Tensor) - Dense Matrix. The shape of the tensor is :math:`(C, M)`. If
+        - **sparse_shape** (tuple(int)) - A positive int tuple which specifies the shape of sparse tensor,
-          `adjoint_dt` is True, its shape must be :math:`(C, M)` after transpose.
+            should have 2 elements, represent sparse tensor shape is :math:`(N, C)`.
+        - **dense** (Tensor) - A 2-D Tensor, the dtype is same as `values`.
+            If `adjoint_st` is False and `adjoint_dt` is False, the shape must be :math:`(C, M)`.
+            If `adjoint_st` is False and `adjoint_dt` is True, the shape must be :math:`(M, C)`.
+            If `adjoint_st` is True and `adjoint_dt` is False, the shape must be :math:`(N, M)`.
+            If `adjoint_st` is True and `adjoint_dt` is True, the shape must be :math:`(M, N)`.
 
     Outputs:
-        Tensor, the shape of tensor  is :math:`(N, M)`. The output data type is the same as "values".
+        Tensor, the dtype is the same as `values`.
+        If `adjoint_st` is False, the shape is :math:`(N, M)`.
+        If `adjoint_st` is True, the shape is :math:`(C, M)`.
 
     Raises:
-        TypeError: If `indices` is neither int32 nor int64.
+        TypeError: If the type of `adjoint_st` or `adjoint_dt` is not bool, or the dtype of `indices`,
-        TypeError: If 'values' is not boot, uint8-64, int8-64, float16-64.
+            dtype of `values` and dtype of `dense` don't meet the parameter description.
-        TypeError: If 'dense' is not boot, uint8-64, int8-64, float16-64.
+        ValueError: If `sparse_shape`, shape of `indices, shape of `values`,
-        ValueError: If length of shape of `SparseTensor` or `DenseTensor` is not equal to 2
+            and shape of `dense` don't meet the parameter description.
 
     Supported Platforms:
         ``CPU``
@@ -93,9 +128,14 @@ class SparseTensorDenseMatmul(PrimitiveWithInfer):
     Examples:
         >>> indices = Tensor([[0, 1], [1, 2]], dtype=ms.int32)
         >>> values = Tensor([1, 2], dtype=ms.float32)
-        >>> dense_shape = (3, 4)
+        >>> sparse_shape = (3, 4)
-        >>> dsMatrix = Tensor([[1,1], [2,2], [3,3 ], [4, 4]], dtype=ms.float32)
+        >>> dense = Tensor([[1,1], [2,2], [3,3 ], [4, 4]], dtype=ms.float32)
-        >>> out = ops.SparseTensorDenseMatmul(indices, values, dense_shape, dsMatrix)
+        >>> sparse_dense_matmul = ops.SparseTensorDenseMatmul()
+        >>> out = sparse_dense_matmul(indices, values, sparse_shape, dense)
+        >>> print(out)
+        [[2 2]
+         [0 6]
+         [6 0]]
     """
 
     @prim_attr_register
@@ -103,27 +143,40 @@ class SparseTensorDenseMatmul(PrimitiveWithInfer):
         """Initialize SparseTensorDenseMatmul"""
         self.adjoint_st = adjoint_st
         self.adjoint_dt = adjoint_dt
-        self.init_prim_io_names(inputs=['indices', 'values', 'dense_shape', 'dense'],
+        self.init_prim_io_names(inputs=['indices', 'values', 'sparse_shape', 'dense'],
                                 outputs=['output'])
         self.add_prim_attr('adjoint_st', self.adjoint_st)
         self.add_prim_attr('adjoint_dt', self.adjoint_dt)
         validator.check_value_type("adjoint_st", adjoint_st, [bool], self.name)
         validator.check_value_type("adjoint_dt", adjoint_dt, [bool], self.name)
 
-    def __infer__(self, indices, values, dense_shape, dense):
+    def __infer__(self, indices, values, sparse_shape, dense):
         validator.check_tensor_dtype_valid('indices', indices['dtype'], [mstype.int32, mstype.int64], self.name)
-        valid_types = mstype.number_type + (mstype.bool_,)
+        valid_types = (mstype.float16, mstype.float32, mstype.float64, mstype.int32, mstype.int64)
         args = {'values': values['dtype'], 'dense': dense['dtype']}
         validator.check_tensors_dtypes_same_and_valid(args, valid_types, self.name)
-        a_shape = dense_shape['value'][::-1] if self.adjoint_st else dense_shape['value']
+        indices_shape = indices['shape']
+        if len(indices_shape) != 2 or indices_shape[1] != 2:
+            raise ValueError("SparseTensorDenseMatmul requires 'indices' must be a 2-D Tensor and "
+                             f"the second dimension length must be 2, but got 'indices' shape: {indices_shape}")
+        values_shape = values['shape']
+        if len(values_shape) != 1 or values_shape[0] != indices_shape[0]:
+            raise ValueError("SparseTensorDenseMatmul requires 'value's must be a 1-D Tensor and "
+                             f"the first dimension length must be equal to the first dimension length of 'indices', "
+                             f"but got 'indices' shape: {indices_shape}, 'values' shape: {values_shape}")
+        a_shape = sparse_shape['value'][::-1] if self.adjoint_st else sparse_shape['value']
         b_shape = dense['shape'][::-1] if self.adjoint_dt else dense['shape']
+        for i in a_shape:
+            if isinstance(i, bool) or not isinstance(i, int) or i <= 0:
+                raise ValueError("SparseTensorDenseMatmul requires all elements in 'sparse_shape' must be "
+                                 f"positive int number, but got sparse shape: {a_shape}")
         if len(a_shape) != 2 or len(b_shape) != 2:
-            raise ValueError('SparseTensorDenseMatmul requires SparseTensor and DenseTensor have the same dimension'
+            raise ValueError("SparseTensorDenseMatmul requires both the 'sparse_shape' length and the dense tensor "
-                             + f'and equal to 2, while SparseTensor dim is ({len(a_shape)}) and DenseTensor dim is '
+                             f"rank should be equal to 2, but got 'sparse_shape' length: {len(a_shape)}, "
-                             + f'({len(b_shape)}).')
+                             f"dense tensor rank: {len(b_shape)}")
         if a_shape[1] != b_shape[0]:
-            raise ValueError('SparseTensorDenseMatmul requires SparseTensor dim_1 should be equal to DenseTensor dim_0,'
+            raise ValueError(f"The sparse tensor shape: {a_shape} and the dense tensor shape: {b_shape} "
-                             f'but got SparseTensor dim_1: {a_shape[1]}, DenseTensor dim_0: {b_shape[0]}')
+                             f"don't meet the condition for matmul")
         out_shape = [a_shape[0], b_shape[1]]
         out = {'shape': tuple(out_shape),
                'dtype': values['dtype'],

tests/st/ops/cpu/test_sparse_tensor_dense_matmul_op.py

@@ -54,10 +54,10 @@ def judge_result_correct(result, expect):
 @pytest.mark.env_onecard
 def test_sparse_tensor_dense_matmul_no_transpose():
     indices_np = np.array([[0, 0], [1, 1], [2, 2], [2, 3]], np.int64)
-    values_np = np.array([2, 3, 4, 5], np.float32)
+    values_np = np.array([2, 3, 4, 5], np.float16)
     dense_shape = (3, 4)
-    sparse_np = np.array([[2, 0, 0, 0], [0, 3, 0, 0], [0, 0, 4, 5]], dtype=np.float32)
+    sparse_np = np.array([[2, 0, 0, 0], [0, 3, 0, 0], [0, 0, 4, 5]], dtype=np.float16)
-    dense_np = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], dtype=np.float32)
+    dense_np = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], dtype=np.float16)
 
     sparse_dense_matmul_net = SparseDenseMatmulNet()
     indices = Tensor(indices_np)
@@ -69,12 +69,12 @@ def test_sparse_tensor_dense_matmul_no_transpose():
 
     grad_net = GradNet(sparse_dense_matmul_net)
     grad_ms = grad_net(indices, values, dense_shape, dense)
-    expect_values_grad = np.array([3., 12., 21., 30.], dtype=np.float32)
+    expect_values_grad = np.array([3., 12., 21., 30.], dtype=np.float16)
     judge_result_correct(grad_ms[1].asnumpy(), expect_values_grad)
     expect_dense_grad = np.array([[2., 2., 2.],
                                   [3., 3., 3.],
                                   [4., 4., 4.],
-                                  [5., 5., 5.]], dtype=np.float32)
+                                  [5., 5., 5.]], dtype=np.float16)
     judge_result_correct(grad_ms[2].asnumpy(), expect_dense_grad)
 
 

tests/st/ops/cpu/test_sparse_to_dense_op.py

@@ -17,182 +17,129 @@ import numpy as np
 import pytest
 import mindspore.context as context
 from mindspore import Tensor
-from mindspore.nn import Cell
+from mindspore import nn
 from mindspore.ops import operations as P
+from mindspore.ops import composite as C
 
-context.set_context(mode=context.GRAPH_MODE, enable_graph_kernel=True, device_target="CPU")
+context.set_context(mode=context.GRAPH_MODE, device_target="CPU")
 
 
-class SparseToDenseNet(Cell):
+class SparseToDenseNet(nn.Cell):
     def __init__(self):
         super(SparseToDenseNet, self).__init__()
         self.sparse_to_dense = P.SparseToDense()
 
-    def construct(self, indices, values, dense_shape):
+    def construct(self, indices, values, sparse_shape):
-        return self.sparse_to_dense(indices, values, dense_shape)
+        return self.sparse_to_dense(indices, values, sparse_shape)
+
+
+class GradNet(nn.Cell):
+    def __init__(self, network):
+        super(GradNet, self).__init__()
+        self.grad = C.GradOperation(get_all=True, sens_param=True)
+        self.network = network
+
+    def construct(self, indices, values, sparse_shape, sense):
+        return self.grad(self.network)(indices, values, sparse_shape, sense)
+
+
+def judge_result_correct(result, expect):
+    assert result.dtype == expect.dtype
+    assert result.shape == expect.shape
+    assert np.allclose(result, expect)
+
+
+def sparse_to_dense_int(i_type, v_type):
+    indices = np.array([[0, 1], [1, 2]]).astype(i_type)
+    values = np.array([7, 8]).astype(v_type)
+    sparse_shape = (3, 3)
+    forward_net = SparseToDenseNet()
+    forward_output = forward_net(Tensor(indices), Tensor(values), sparse_shape)
+    expect_forward_output = np.array([[0, 7, 0],
+                                      [0, 0, 8],
+                                      [0, 0, 0]]).astype(v_type)
+    judge_result_correct(forward_output.asnumpy(), expect_forward_output)
+
+    grad_net = GradNet(forward_net)
+    sense = Tensor(np.arange(9).reshape((3, 3)).astype(v_type))
+    grad_output = grad_net(Tensor(indices), Tensor(values), sparse_shape, sense)
+    expect_grad_output = np.array([1, 5]).astype(v_type)
+    judge_result_correct(grad_output[1].asnumpy(), expect_grad_output)
+
+
+def sparse_to_dense_float(i_type, v_type):
+    indices = np.array([[0, 1, 0], [1, 2, 1], [2, 3, 2], [0, 2, 3]]).astype(i_type)
+    values = np.array([6.5, 7.5, 9.5, 10.5]).astype(v_type)
+    sparse_shape = (3, 4, 4)
+    forward_net = SparseToDenseNet()
+    forward_output = forward_net(Tensor(indices), Tensor(values), sparse_shape)
+    expect_forward_output = np.array([[[0, 0, 0, 0],
+                                       [6.5, 0, 0, 0],
+                                       [0, 0, 0, 10.5],
+                                       [0, 0, 0, 0]],
+                                      [[0, 0, 0, 0],
+                                       [0, 0, 0, 0],
+                                       [0, 7.5, 0, 0],
+                                       [0, 0, 0, 0]],
+                                      [[0, 0, 0, 0],
+                                       [0, 0, 0, 0],
+                                       [0, 0, 0, 0],
+                                       [0, 0, 9.5, 0]]]).astype(v_type)
+    judge_result_correct(forward_output.asnumpy(), expect_forward_output)
+
+    grad_net = GradNet(forward_net)
+    sense = Tensor(np.arange(48).reshape((3, 4, 4)).astype(v_type) + 0.8)
+    grad_output = grad_net(Tensor(indices), Tensor(values), sparse_shape, sense)
+    expect_grad_output = np.array([4.8, 25.8, 46.8, 11.8]).astype(v_type)
+    judge_result_correct(grad_output[1].asnumpy(), expect_grad_output)
+
+
+def sparse_to_dense_1D(i_type, v_type):
+    indices = np.array([[8], [2], [6], [4]]).astype(i_type)
+    values = np.array([6.5, 7.5, 9.5, 10.5]).astype(v_type)
+    sparse_shape = (10,)
+    forward_net = SparseToDenseNet()
+    forward_output = forward_net(Tensor(indices), Tensor(values), sparse_shape)
+    expect_forward_output = np.array([0, 0, 7.5, 0, 10.5, 0, 9.5, 0, 6.5, 0]).astype(v_type)
+    judge_result_correct(forward_output.asnumpy(), expect_forward_output)
+
+    grad_net = GradNet(forward_net)
+    sense = Tensor(np.arange(10).astype(v_type) + 0.8)
+    grad_output = grad_net(Tensor(indices), Tensor(values), sparse_shape, sense)
+    expect_grad_output = np.array([8.8, 2.8, 6.8, 4.8]).astype(v_type)
+    judge_result_correct(grad_output[1].asnumpy(), expect_grad_output)
+
+
+indices_types = (np.int32, np.int64)
 
 
 @pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
+@pytest.mark.platform_x86_cpu
 @pytest.mark.env_onecard
-def test_sparse_to_dense_A():
+def test_sparse_to_dense_int():
-    np.random.seed(0)
+    values_types = (np.bool_,
-    indices = np.array([[0, 1], [1, 2]]).astype(np.int32)
+                    np.uint8, np.uint16, np.uint32, np.uint64,
-    values = np.array([7, 8]).astype(np.int32)
+                    np.int8, np.int16, np.int32, np.int64)
-    dense_shape = (3, 4)
+    for i_type in indices_types:
-    net = SparseToDenseNet()
+        for v_type in values_types:
-    result = net(Tensor(indices), Tensor(values), dense_shape)
+            sparse_to_dense_int(i_type, v_type)
-    expect_output = np.array([[0, 7, 0, 0],
-                              [0, 0, 8, 0],
-                              [0, 0, 0, 0]]).astype(np.int32)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
 
 
 @pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
+@pytest.mark.platform_x86_cpu
 @pytest.mark.env_onecard
-def test_sparse_to_dense_B():
+def test_sparse_to_dense_float():
-    np.random.seed(0)
+    values_types = (np.float16, np.float32, np.float64)
-    indices = np.array([[0, 1], [1, 2], [2, 3]]).astype(np.int32)
+    for i_type in indices_types:
-    values = np.array([6.5, 7.5, 9.5]).astype(np.float64)
+        for v_type in values_types:
-    dense_shape = (3, 4)
+            sparse_to_dense_float(i_type, v_type)
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    expect_output = np.array([[0, 6.5, 0, 0],
-                              [0, 0, 7.5, 0],
-                              [0, 0, 0, 9.5]]).astype(np.float64)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
 
 
 @pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
+@pytest.mark.platform_x86_cpu
 @pytest.mark.env_onecard
-def test_sparse_to_dense_C():
+def test_sparse_to_dense_1D():
-    np.random.seed(0)
+    values_types = (np.float16, np.float32, np.float64)
-    indices = np.array([[0, 1, 0, 0],
+    for i_type in indices_types:
-                        [1, 0, 0, 2],
+        for v_type in values_types:
-                        [2, 0, 3, 0],
+            sparse_to_dense_1D(i_type, v_type)
-                        [4, 2, 3, 5]]).astype(np.int32)
-    values = np.array([26.5, 17.5, 39.5, 11.5]).astype(np.float16)
-    dense_shape = (10, 8, 5, 10)
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    expect_output = np.zeros([10, 8, 5, 10]).astype(np.float16)
-    for i in range(0, indices.shape[0]):
-        j = indices[i][0]
-        k = indices[i][1]
-        l = indices[i][2]
-        m = indices[i][3]
-        expect_output[j][k][l][m] = values[i]
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
-
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
-@pytest.mark.env_onecard
-def test_sparse_to_dense_D():
-    np.random.seed(0)
-    indices = np.array([[0, 1, 0, 0, 2, 1],
-                        [9, 0, 0, 8, 0, 0],
-                        [2, 0, 4, 0, 1, 1],
-                        [4, 2, 3, 5, 0, 2],
-                        [7, 4, 3, 9, 0, 1]]).astype(np.int32)
-    values = np.array([1, 1, 1, 1, 1]).astype(np.bool)
-    dense_shape = (10, 5, 5, 10, 3, 3)
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    expect_output = np.zeros([10, 5, 5, 10, 3, 3]).astype(np.bool)
-    for i in range(0, indices.shape[0]):
-        j = indices[i][0]
-        k = indices[i][1]
-        l = indices[i][2]
-        m = indices[i][3]
-        u = indices[i][4]
-        v = indices[i][5]
-        expect_output[j][k][l][m][u][v] = values[i]
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
-
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
-@pytest.mark.env_onecard
-def test_sparse_to_dense_E():
-    indices = np.array([2, 5, 7]).astype(np.int32)
-    values = np.array([17, 18, 19]).astype(np.int8)
-    dense_shape = ([10])
-    expect_output = np.zeros([10]).astype(np.int8)
-    for i in range(0, indices.shape[0]):
-        j = indices[i]
-        expect_output[j] = values[i]
-
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
-
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
-@pytest.mark.env_onecard
-def test_sparse_to_dense_F():
-    indices = np.array([2, 4, 18]).astype(np.int32)
-    values = np.array([-23, 18, -1]).astype(np.int16)
-    dense_shape = ([20])
-    expect_output = np.zeros([20]).astype(np.int16)
-    for i in range(0, indices.shape[0]):
-        j = indices[i]
-        expect_output[j] = values[i]
-
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
-
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
-@pytest.mark.env_onecard
-def test_sparse_to_dense_G():
-    indices = np.array([2, 5, 7]).astype(np.int32)
-    values = np.array([17, 18, 19]).astype(np.uint8)
-    dense_shape = ([10])
-    expect_output = np.zeros([10]).astype(np.uint8)
-    for i in range(0, indices.shape[0]):
-        j = indices[i]
-        expect_output[j] = values[i]
-
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
-
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
-@pytest.mark.env_onecard
-def test_sparse_to_dense_H():
-    indices = np.array([2, 5, 7]).astype(np.int32)
-    values = np.array([17, 18, 19]).astype(np.uint16)
-    dense_shape = ([10])
-    expect_output = np.zeros([10]).astype(np.uint16)
-    for i in range(0, indices.shape[0]):
-        j = indices[i]
-        expect_output[j] = values[i]
-
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)
-
-
-@pytest.mark.level0
-@pytest.mark.platform_x86_cpu_training
-@pytest.mark.env_onecard
-def test_sparse_to_dense_I():
-    indices = np.array([2, 5, 7]).astype(np.int64)
-    values = np.array([17, 18, 19]).astype(np.float16)
-    dense_shape = ([10])
-    expect_output = np.zeros([10]).astype(np.float16)
-    for i in range(0, indices.shape[0]):
-        j = indices[i]
-        expect_output[j] = values[i]
-
-    net = SparseToDenseNet()
-    result = net(Tensor(indices), Tensor(values), dense_shape)
-    assert np.allclose(result.asnumpy(), expect_output, rtol=1.e-4, atol=1.e-8, equal_nan=True)