gather modify

paddle/operators/gather_func.h

@@ -21,20 +21,17 @@ limitations under the License. */
 /**
  * Return a new tensor from source tensor, gathered according to index
  * input[src]: type-T source Tensor
- * input[Index]: type-int index Tensor (1-D)
+ * input[index]: type-int index Tensor (1-D)
  * return: output tensor
  */
-template <typename place, typename T>
-Tensor* Gather_func(Tensor* Src, Tensor* Index) {
-	// assert index is an int-type tensor?
-	// assert(Index->istype(int));
-
+template <typename Place, typename T>
+Tensor* Gather(Tensor* src, Tensor* index) {
   // check index of shape 1-D
-	assert(Index->dims().size()==1);
-	int index_size = Index->dims()[0];
+  PADDLE_ENFORCE(index->dims().size()==1);
+  int index_size = index->dims()[0];
 
   // Source shape
-	auto src_dims = Src->dims();
+  auto src_dims = src->dims();
   DDim output_dims(dims_src);
   // Create a tensor of shape [index_size, dim_src[1:]]
   output_dims[0] = index_size;
@@ -51,12 +48,12 @@ Tensor* Gather_func(Tensor* Src, Tensor* Index) {
   if (place == CPUPlace()) {
 	// init for CPU
 	output = New_tensor.mutable_data<T>(output_dims, CPUPlace());
-		CPUGather(Src->data(), Index->data(), slice_size, new_tensor->mutable_data());
+	CPUGather(src->data(), index->data(), slice_size, new_tensor->mutable_data());
   } else { // GPU
 	// init for GPU
 	output = New_tensor.mutable_data<T>(output_dims, GPUPlace());
 	/* how to specialize device??*/
-		GPUGather(d, Src->data(), Index->data(), slice_size, new_tensor->mutable_data());
+	GPUGather(d, src->data(), index->data(), slice_size, new_tensor->mutable_data());
   }
   return New_tensor;
 }
@@ -82,7 +79,7 @@ void CPUGather(const T* params, const int* indices,
 */
 template<typename T>
 void GPUGather(const GPUDevice& d,
-			   const T* src, const int* Index, 
+			   const T* src, const int* index, 
 	           const int slice_size, const int index_size,
 	           T* output) {
   int block_count = slice_size * index_size;
@@ -90,7 +87,7 @@ void GPUGather(const GPUDevice& d,
 
   GatherOpKernel<T>
           <<<block_count, thread_per_block, 0, d.stream()>>>(
-              src, Index, output, slice_size,
+              src, index, output, slice_size,
               indices_size, slice_size, out_size);
 }