Revert "[OP] take (#4538)" (#4627)

This reverts commit 8159ad6.

src/operator/tensor/indexing_op.cc

@@ -1,8 +1,8 @@
 /*!
- * Copyright (c) 2017 by Contributors
+ * Copyright (c) 2016 by Contributors
  * \file indexing_op.cc
  * \brief
- * \author Siyi Li, Chi Zhang
+ * \author Siyi Li
 */
 
 #include "./indexing_op.h"
@@ -28,57 +28,25 @@ NNVM_REGISTER_OP(Embedding)
   [](const NodeAttrs& attrs) {
     return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
   })
-.set_attr<FCompute>("FCompute<cpu>", TakeOpForward<cpu>)
+.set_attr<FCompute>("FCompute<cpu>", EmbeddingOpForward<cpu>)
 .set_attr<nnvm::FGradient>("FGradient",
   [](const nnvm::NodePtr& n,  const std::vector<nnvm::NodeEntry>& ograds) {
     std::vector<nnvm::NodeEntry> heads(ograds.begin(), ograds.end());
     heads.push_back(n->inputs[0]);
-    return MakeGradNode("_backward_take", n, heads, n->attrs.dict);
+    return MakeGradNode("_backward_Embedding", n, heads, n->attrs.dict);
   })
 .add_argument("data", "Symbol", "Input data to the EmbeddingOp.")
 .add_argument("weight", "Symbol", "Embedding weight matrix.")
 .add_arguments(EmbeddingParam::__FIELDS__());
 
-DMLC_REGISTER_PARAMETER(TakeParam);
-
-NNVM_REGISTER_OP(take)
-.MXNET_DESCRIBE("Take row vectors from a 2D matrix according to the indices"
-                " For an input of index with shape (d1, ..., dK), the output"
-                " shape is (d1, ..., dK, row_vector_length).All the input"
-                " values should be integers in the range"
-                " [0, column_vector_length).")
-.set_num_inputs(2)
-.set_num_outputs(1)
-.set_attr_parser(ParamParser<TakeParam>)
-.set_attr<nnvm::FListInputNames>("FListInputNames",
-  [](const NodeAttrs& attrs) {
-    return std::vector<std::string>{"idx", "data"};
-  })
-.set_attr<nnvm::FInferShape>("FInferShape", TakeOpShape)
-.set_attr<nnvm::FInferType>("FInferType", TakeOpType)
-.set_attr<FResourceRequest>("FResourceRequest",
-  [](const NodeAttrs& attrs) {
-    return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
-  })
-.set_attr<FCompute>("FCompute<cpu>", TakeOpForward<cpu>)
-.set_attr<nnvm::FGradient>("FGradient",
-  [](const nnvm::NodePtr& n,  const std::vector<nnvm::NodeEntry>& ograds) {
-    std::vector<nnvm::NodeEntry> heads(ograds.begin(), ograds.end());
-    heads.push_back(n->inputs[0]);
-    return MakeGradNode("_backward_take", n, heads, n->attrs.dict);
-  })
-.add_argument("idx", "Symbol", "Tensor that records the indices to be taken in data.")
-.add_argument("data", "Symbol", "2D matrix to be taken.")
-.add_arguments(TakeParam::__FIELDS__());
-
-NNVM_REGISTER_OP(_backward_take)
+NNVM_REGISTER_OP(_backward_Embedding)
 .set_num_inputs(2)
 .set_num_outputs(2)
 .set_attr<FResourceRequest>("FResourceRequest",
   [](const NodeAttrs& attrs) {
     return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
   })
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
-.set_attr<FCompute>("FCompute<cpu>", TakeOpBackward<cpu>);
+.set_attr<FCompute>("FCompute<cpu>", EmbeddingOpBackward<cpu>);
 }  // namespace op
 }  // namespace mxnet

src/operator/tensor/indexing_op.cu

@@ -1,21 +1,18 @@
 /*!
- * Copyright (c) 2017 by Contributors
+ * Copyright (c) 2016 by Contributors
  * \file indexing_op.cu
  * \brief
- * \author Siyi Li, Chi Zhang
+ * \author Siyi Li
 */
 
 #include "./indexing_op.h"
 namespace mxnet {
 namespace op {
 NNVM_REGISTER_OP(Embedding)
-.set_attr<FCompute>("FCompute<gpu>", TakeOpForward<gpu>);
+.set_attr<FCompute>("FCompute<gpu>", EmbeddingOpForward<gpu>);
 
-NNVM_REGISTER_OP(take)
-.set_attr<FCompute>("FCompute<gpu>", TakeOpForward<gpu>);
-
-NNVM_REGISTER_OP(_backward_take)
-.set_attr<FCompute>("FCompute<gpu>", TakeOpBackward<gpu>);
+NNVM_REGISTER_OP(_backward_Embedding)
+.set_attr<FCompute>("FCompute<gpu>", EmbeddingOpBackward<gpu>);
 }  // namespace op
 }  // namespace mxnet
 

src/operator/tensor/indexing_op.h

@@ -1,8 +1,8 @@
 /*!
- * Copyright (c) 2017 by Contributors
+ * Copyright (c) 2016 by Contributors
  * \file indexing_op.h
  * \brief
- * \author Bing Xu, Siyi Li, Chi Zhang
+ * \author Bing Xu, Siyi Li
 */
 #ifndef MXNET_OPERATOR_TENSOR_INDEXING_OP_H_
 #define MXNET_OPERATOR_TENSOR_INDEXING_OP_H_
@@ -22,12 +22,11 @@
 namespace mxnet {
 namespace op {
 
-namespace indexing {
-enum IndexingOpInputs {kIdx, kData};
-enum IndexingOpOutputs {kOut};
-enum IndexingOpResource {kTempSpace};
-enum IndexingOpMode {kRaise, kWrap, kClip};
-}  // namespace indexing
+namespace embedding {
+enum EmbeddingOpInputs {kData, kWeight};
+enum EmbeddingOpOutputs {kOut};
+enum EmbeddingOpResource {kTempSpace};
+}  // namespace embedding
 
 struct EmbeddingParam: public dmlc::Parameter<EmbeddingParam> {
   int input_dim;
@@ -44,11 +43,11 @@ inline bool EmbeddingOpShape(const nnvm::NodeAttrs& attrs,
                              std::vector<TShape> *in_attrs,
                              std::vector<TShape> *out_attrs) {
   using namespace mshadow;
-  const TShape &dshape = (*in_attrs)[indexing::kIdx];
+  const TShape &dshape = (*in_attrs)[embedding::kData];
   if (dshape.ndim() ==  0) return false;
   const EmbeddingParam& param = nnvm::get<EmbeddingParam>(attrs.parsed);
-  SHAPE_ASSIGN_CHECK(*in_attrs, indexing::kData, Shape2(param.input_dim,
-                                                        param.output_dim));
+  SHAPE_ASSIGN_CHECK(*in_attrs, embedding::kWeight, Shape2(param.input_dim,
+                                                           param.output_dim));
   out_attrs->clear();
 
   TShape oshape(dshape.ndim()+1);
@@ -81,154 +80,84 @@ inline bool EmbeddingOpType(const nnvm::NodeAttrs& attrs,
   return true;
 }
 
-// TODO(somebody): behaviors specified by params
-struct TakeParam: public dmlc::Parameter<TakeParam> {
-  int axis;
-  int mode;
-  DMLC_DECLARE_PARAMETER(TakeParam) {
-    DMLC_DECLARE_FIELD(axis)
-    .set_lower_bound(0)
-    .set_default(0)
-    .describe("the axis of data tensor to be taken.");
-    DMLC_DECLARE_FIELD(mode)
-    .add_enum("raise", indexing::kRaise)
-    .add_enum("wrap", indexing::kWrap)
-    .add_enum("clip", indexing::kClip)
-    .set_default(indexing::kRaise)
-    .describe("specify how out-of-bound indices bahave.");
-  }
-};
-
-inline bool TakeOpShape(const nnvm::NodeAttrs& attrs,
-                        std::vector<TShape> *in_attrs,
-                        std::vector<TShape> *out_attrs) {
-    using namespace mshadow;
-    const TShape &arrshape = (*in_attrs)[indexing::kData];
-    const TShape &idxshape = (*in_attrs)[indexing::kIdx];
-    if (idxshape.ndim() == 0) return false;
-    out_attrs->clear();
-
-    TShape oshape(idxshape.ndim() + arrshape.ndim() - 1);
-    for (size_t i = 0; i < idxshape.ndim(); ++i) {
-        oshape[i] = idxshape[i];
-    }
-    for (size_t i = 0; i < arrshape.ndim() - 1; i++) {
-        oshape[i + idxshape.ndim()] = arrshape[i + 1];
-    }
-    out_attrs->push_back(oshape);
-    return true;
+template<typename xpu>
+void EmbeddingOpForward(const nnvm::NodeAttrs& attrs,
+                        const OpContext& ctx,
+                        const std::vector<TBlob>& inputs,
+                        const std::vector<OpReqType>& req,
+                        const std::vector<TBlob>& outputs) {
+  using namespace mshadow;
+  using namespace mshadow::expr;
+  CHECK_EQ(req[embedding::kOut], kWriteTo);
+  CHECK_EQ(inputs.size(), 2);
+  CHECK_EQ(outputs.size(), 1);
+  CHECK_EQ(inputs[embedding::kWeight].ndim(), 2)
+          << "Embedding layer expects its weight to be two-dimensional. "
+          << inputs[embedding::kWeight].ndim()
+          << " dimensional input is given instead";
+
+  const TShape& ishape = inputs[embedding::kData].shape_;
+  const TShape& oshape = outputs[embedding::kOut].shape_;
+
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+    Tensor<xpu, 1, DType> data = inputs[embedding::kData].get_with_shape<xpu, 1, DType>(
+      Shape1(ishape.ProdShape(0, ishape.ndim())), s);
+    Tensor<xpu, 2, DType> wmat = inputs[embedding::kWeight].get<xpu, 2, DType>(s);
+    Tensor<xpu, 2, DType> out = outputs[embedding::kOut].get_with_shape<xpu, 2, DType>(
+      Shape2(oshape.ProdShape(0, oshape.ndim()-1), oshape[oshape.ndim()-1]), s);
+    out = take(data, wmat);
+  });
 }
 
-inline bool TakeOpType(const nnvm::NodeAttrs& attrs,
-                       std::vector<int> *in_type,
-                       std::vector<int> *out_type) {
-  // using single dtype ("float32") for safety reason
-  CHECK_GE(in_type->size(), 2);
-  int dtype = (*in_type)[1];
-  CHECK_NE(dtype, -1) << "idx must have specified type";
-  for (index_t i = 0; i < in_type->size(); ++i) {
-    if ((*in_type)[i] == -1) {
-      (*in_type)[i] = dtype;
+template<typename xpu>
+void EmbeddingOpBackward(const nnvm::NodeAttrs& attrs,
+                         const OpContext& ctx,
+                         const std::vector<TBlob>& inputs,
+                         const std::vector<OpReqType>& req,
+                         const std::vector<TBlob>& outputs) {
+  using namespace mshadow;
+  using namespace mshadow::expr;
+  CHECK_EQ(inputs.size(), 2);
+  CHECK_EQ(outputs.size(), 2);
+  CHECK_EQ(req[embedding::kData], kNullOp)
+          << "Embedding layer doesn't support calculate data gradient";
+
+  const TShape& ishape = inputs[1].shape_;
+  const TShape& oshape = inputs[0].shape_;
+
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+    Tensor < xpu, 1, DType > data = inputs[1].get_with_shape<xpu, 1, DType>(
+      Shape1(ishape.ProdShape(0, ishape.ndim())), s);
+    Tensor<xpu, 2, DType> grad_out = inputs[0].get_with_shape<xpu, 2, DType>(
+    Shape2(oshape.ProdShape(0, oshape.ndim()-1), oshape[oshape.ndim()-1]), s);
+    Tensor<xpu, 2, DType> grad_in = outputs[1].get<xpu, 2, DType>(s);
+
+
+    if (req[embedding::kWeight] == kWriteTo || req[embedding::kWeight] == kAddTo) {
+      if (req[embedding::kWeight] == kWriteTo) {
+        grad_in = scalar<DType>(0.0f);
+      }
+      if ((grad_out.shape_[0] < grad_out.shape_[1]) && (grad_out.shape_[0] < 512)) {
+        AddTakeGrad(grad_in, data, grad_out);
+      } else {
+        Tensor<xpu, 2, int> workspace =
+          ctx.requested[embedding::kTempSpace].get_space_typed<xpu, 2, int>(
+            mshadow::Shape2(2, data.shape_.Size()), s);
+        Tensor<xpu, 1, int> sorted_data = workspace[0];
+        Tensor<xpu, 1, int> original_index = workspace[1];
+        sorted_data = tcast<int>(data);
+        original_index = range<int>(0, data.shape_.Size());
+        SortByKey(sorted_data, original_index, true);
+        AddTakeGradLargeBatch(grad_in, sorted_data, original_index, grad_out);
+      }
     } else {
-      CHECK_EQ((*in_type)[i], dtype) << "This layer requires uniform type. "
-                                     << "Expected " << dtype << " v.s. given "
-                                     << (*in_type)[i];
+      LOG(FATAL) << "wrong req";
     }
-  }
-  out_type->clear();
-  out_type->push_back(dtype);
-  return true;
-}
-
-template<typename xpu>
-void TakeOpForward(const nnvm::NodeAttrs& attrs,
-                   const OpContext& ctx,
-                   const std::vector<TBlob>& inputs,
-                   const std::vector<OpReqType>& req,
-                   const std::vector<TBlob>& outputs) {
-    using namespace mshadow;
-    using namespace mshadow::expr;
-    CHECK_EQ(req[indexing::kOut], kWriteTo);
-    CHECK_EQ(inputs.size(), 2);
-    CHECK_EQ(outputs.size(), 1);
-    CHECK_GE(inputs[indexing::kData].ndim(), 2)
-        << "Indexing layer expects its array's size to be at least 2. "
-        << inputs[indexing::kData].ndim()
-        << " dimensional input is given instead";
-
-    const TShape& idxshape = inputs[indexing::kIdx].shape_;
-    const TShape& arrshape = inputs[indexing::kData].shape_;
-    const TShape& oshape = outputs[indexing::kOut].shape_;
-
-    int idxndim = idxshape.ndim();
-
-    Stream<xpu> *s = ctx.get_stream<xpu>();
-    MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
-        Tensor<xpu, 1, DType> idx = inputs[indexing::kIdx].get_with_shape<xpu, 1, DType>(
-            Shape1(idxshape.ProdShape(0, idxndim)), s);
-        Tensor<xpu, 2, DType> data = inputs[indexing::kData].get_with_shape<xpu, 2, DType>(
-            Shape2(arrshape[0], arrshape.ProdShape(1, arrshape.ndim())), s);
-        Tensor<xpu, 2, DType> out = outputs[indexing::kOut].get_with_shape<xpu, 2, DType>(
-            Shape2(oshape.ProdShape(0, idxndim), oshape.ProdShape(idxndim, oshape.ndim())), s);
-        out = take(idx, data);
-    });
+  });
 }
 
-template<typename xpu>
-void TakeOpBackward(const nnvm::NodeAttrs& attrs,
-                    const OpContext& ctx,
-                    const std::vector<TBlob>& inputs,
-                    const std::vector<OpReqType>& req,
-                    const std::vector<TBlob>& outputs) {
-    using namespace mshadow;
-    using namespace mshadow::expr;
-    CHECK_EQ(inputs.size(), 2);
-    CHECK_EQ(outputs.size(), 2);
-    CHECK_EQ(req[indexing::kIdx], kNullOp)
-        << "Indexing op doesn't support gradient into index";
-
-    // inputs are specified in the .cc file, which are the gradients from
-    // the upper layer and the input index
-    // outputs are the gradients of inputs in the feed-forward pass
-    const TShape& idxshape = inputs[1].shape_;
-    const TShape& arrshape = outputs[1].shape_;
-    const TShape& oshape = inputs[0].shape_;
-
-    int idxndim = idxshape.ndim();
-
-    // grad_out is the gradient of the outputs in the feed-forward
-    // grad_in is the gradient of the inputs in the feed-forward
-    Stream<xpu> *s = ctx.get_stream<xpu>();
-    MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
-        Tensor<xpu, 1, DType> idx = inputs[1].get_with_shape<xpu, 1, DType>(
-            Shape1(idxshape.ProdShape(0, idxndim)), s);
-        Tensor<xpu, 2, DType> grad_out = inputs[0].get_with_shape<xpu, 2, DType>(
-            Shape2(oshape.ProdShape(0, idxndim), oshape.ProdShape(idxndim, oshape.ndim())), s);
-        Tensor<xpu, 2, DType> grad_in = outputs[1].get_with_shape<xpu, 2, DType>(
-            Shape2(arrshape[0], arrshape.ProdShape(1, arrshape.ndim())), s);
-
-        if (req[indexing::kData] == kWriteTo || req[indexing::kData] == kAddTo) {
-            if (req[indexing::kData] == kWriteTo) {
-                grad_in = scalar<DType>(0.0f);
-            }
-            if ((grad_out.shape_[0] < grad_out.shape_[1]) && (grad_out.shape_[0] < 512)) {
-                AddTakeGrad(grad_in, idx, grad_out);
-            } else {
-                Tensor<xpu, 2, int> workspace =
-                    ctx.requested[indexing::kTempSpace].get_space_typed<xpu, 2, int>(
-                        mshadow::Shape2(2, idx.shape_.Size()), s);
-                Tensor<xpu, 1, int> sorted_idx = workspace[0];
-                Tensor<xpu, 1, int> original_idx = workspace[1];
-                sorted_idx = tcast<int>(idx);
-                original_idx = range<int>(0, idx.shape_.Size());
-                SortByKey(sorted_idx, original_idx, true);
-                AddTakeGradLargeBatch(grad_in, sorted_idx, original_idx, grad_out);
-            }
-        } else {
-            LOG(FATAL) << "wrong req";
-        }
-    });
-}
 
 }  // namespace op
 }  // namespace mxnet