add as an operator

paddle/operators/CMakeLists.txt

@@ -43,6 +43,8 @@ endfunction()
 
 add_subdirectory(math)
 cc_test(gather_test SRCS gather_test.cc DEPS tensor)
+cc_library(gather_op SRCS gather_op.cc DEPS op_registry)
+# cc_test(gather_op_test SRCS gather_op_test.cc DEPS gather_op)
 
 cc_test(scatter_test SRCS scatter_test.cc DEPS tensor)
 

paddle/operators/gather_op.cc

@@ -0,0 +1,64 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/operators/gather_op.h"
+#include "paddle/framework/ddim.h"
+
+namespace paddle {
+namespace operators {
+
+class GatherOp : public framework::OperatorWithKernel {
+ protected:
+  void InferShape(const framework::InferShapeContext &ctx) const override {
+    PADDLE_ENFORCE(ctx.InputSize() == 2, "");
+    PADDLE_ENFORCE(ctx.OutputSize() == 1, "");
+    int batch_size = ctx.Input<Tensor>(1)->dims()[0];
+    PADDLE_ENFORCE(batch_size > 0);
+  }
+};
+
+class GatherOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  GatherOpMaker(framework::OpProto *proto, framework::OpAttrChecker *op_checker)
+      : OpProtoAndCheckerMaker(proto, op_checker) {
+    AddInput("X", "The source input of gather op");
+    AddInput("Index", "The index input of gather op");
+    AddOutput("Y", "The output of add op");
+    AddComment(R"DOC(
+Gather Operator by selecting from the first axis, 
+
+Y = X[Index]
+)DOC");
+  }
+};
+
+class GatherGradOp : public framework::OperatorWithKernel {
+ protected:
+  void InferShape(const framework::InferShapeContext &ctx) const override {
+    // ctx.Output<Tensor>("X" + framework::kGradVarSuffix)
+    //    ->Resize(ctx.Input<Tensor>("X")->dims());
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OP(gather, ops::GatherOp, ops::GatherOpMaker);
+REGISTER_OP_CPU_KERNEL(gather,
+                       ops::GatherOpKernel<paddle::platform::CPUPlace, float>);
+REGISTER_GRADIENT_OP(gather, gather_grad, ops::GatherGradOp);
+REGISTER_OP_CPU_KERNEL(
+    gather_grad,
+    ops::GatherGradientOpKernel<paddle::platform::CPUPlace, float>);

paddle/operators/gather_op.h

@@ -0,0 +1,52 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+#include "gather.h"
+#include "paddle/framework/eigen.h"
+#include "paddle/framework/op_registry.h"
+#include "scatter.h"
+
+namespace paddle {
+namespace operators {
+
+using Tensor = framework::Tensor;
+
+template <typename Place, typename T>
+class GatherOpKernel : public framework::OpKernel {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    auto X = ctx.Input<Tensor>("X");
+    auto Index = ctx.Input<Tensor>("Index");
+    auto Y = ctx.Output<Tensor>("Y");
+
+    Y->mutable_data<T>(ctx.GetPlace());
+    Gather<T>(ctx.GetPlace(), X, Index, Y);
+  }
+};
+
+template <typename Place, typename T>
+class GatherGradientOpKernel : public framework::OpKernel {
+ public:
+  void Compute(const framework::ExecutionContext& ctx) const override {
+    auto Index = ctx.Input<Tensor>("Index");
+    auto dX = ctx.Output<Tensor>(framework::GradVarName("X"));
+    auto dY = ctx.Input<Tensor>(framework::GradVarName("Y"));
+
+    ScatterUpdate<T>(ctx.GetPlace(), dY, Index, dX);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle