Add index_copy() operator

src/operator/contrib/index_copy-inl.h

@@ -0,0 +1,121 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+ /*!
+  * \file index_copy-inl.h
+  * \brief implementation of index_copy tensor operation
+  */
+
+#ifndef MXNET_OPERATOR_CONTRIB_INDEX_COPY_INL_H_
+#define MXNET_OPERATOR_CONTRIB_INDEX_COPY_INL_H_
+
+#include <mxnet/operator_util.h>
+#include <vector>
+#include <limits>
+#include "../elemwise_op_common.h"
+#include "../mshadow_op.h"
+#include "../mxnet_op.h"
+
+namespace mxnet {
+namespace op {
+
+struct copy_all {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i,
+                                  DType* old_tensor,
+                                  DType* out_tensor) {
+    *(out_tensor + i) = *(old_tensor + i);
+  }
+};
+
+struct index_copy {
+  template<typename DType, typename IType>
+  MSHADOW_XINLINE static void Map(int i,
+                                  int dim,
+                                  IType* index,
+                                  DType* new_tensor,
+                                  DType* out_tensor) {
+    DType* out_ptr = out_tensor + static_cast<int>(index[i]) * dim;
+    DType* new_ptr = new_tensor + i * dim;
+    for (int idx = 0; idx < dim; ++idx) {
+      *(out_ptr + idx) = *(new_ptr + idx);
+    }
+  }
+};
+
+template<typename xpu>
+void IndexCopyCompute(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 mxnet_op;
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+  const TBlob& out = outputs[0];
+  const TBlob& idx = inputs[1];
+  int dim = inputs[2].Size() / inputs[1].Size();
+  // copy all
+  MSHADOW_TYPE_SWITCH(out.type_flag_, DType, {
+    Kernel<copy_all, xpu>::Launch(s, inputs[0].Size(),
+                            inputs[0].dptr<DType>(),
+                            outputs[0].dptr<DType>());
+  })
+  // index copy
+  MSHADOW_TYPE_SWITCH(out.type_flag_, DType, {
+    MSHADOW_TYPE_SWITCH(idx.type_flag_, IType, {
+      Kernel<index_copy, xpu>::Launch(s, inputs[1].Size(),
+                            dim,
+                            inputs[1].dptr<IType>(),
+                            inputs[2].dptr<DType>(),
+                            outputs[0].dptr<DType>());
+    })
+  })
+}
+
+inline bool IndexCopyShape(const nnvm::NodeAttrs& attrs,
+                           std::vector<TShape> *in_attrs,
+                           std::vector<TShape> *out_attrs) {
+  CHECK_EQ(in_attrs->size(), 3U);
+  CHECK_EQ(out_attrs->size(), 1U);
+  CHECK_EQ(in_attrs->at(1).ndim(), 1);
+  CHECK_EQ(in_attrs->at(0).ndim(), in_attrs->at(2).ndim());
+  for (size_t i = 0; i < in_attrs->at(0).ndim(); ++i) {
+    if (i == 0) {
+      CHECK_GE(in_attrs->at(0)[i], in_attrs->at(2)[i]);
+    } else {
+      CHECK_EQ(in_attrs->at(0)[i], in_attrs->at(2)[i]);
+    }
+  }
+  CHECK_EQ(in_attrs->at(1)[0], in_attrs->at(2)[0]);
+  SHAPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+  return true;
+}
+
+inline bool IndexCopyType(const nnvm::NodeAttrs& attrs,
+                          std::vector<int> *in_attrs,
+                          std::vector<int> *out_attrs) {
+  TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
+  return true;
+}
+
+}  // namespace op
+}  // namespace mxnet
+
+#endif  // MXNET_OPERATOR_CONTRIB_INDEX_COPY_INL_H_

src/operator/contrib/index_copy.cc

@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * \file index_copy.cc
+ * \brief
+ */
+#include "./index_copy-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_contrib_index_copy)
+.describe(R"code(Copies the elements of a `new_tensor` into the `old_tensor` by 
+selecting the indices in the order given in `index`. The output will be a new tensor 
+contains the rest elements of old tensor and the copied elements of new tensor. 
+For example, if `index[i] == j`, then the `i`th row of `new_tensor` is copied to the 
+`j`th row of output.
+
+The `index` must be a vector and it must have the same size with the `0`th dimimention of 
+`new_tensor`. Also, the `0`th dimimention of old_tensor must `>=` the `0`th dimimention of 
+`new_tensor`, or an error will be raised.
+
+Examples::
+
+x = mx.nd.zeros((5,3))
+t = mx.nd.array([[1,2,3],[4,5,6],[7,8,9]])
+index = mx.nd.array([0,4,2])
+
+mx.nd.contrib.index_copy(x, index, t)
+
+[[1. 2. 3.]
+ [0. 0. 0.]
+ [7. 8. 9.]
+ [0. 0. 0.]
+ [4. 5. 6.]]
+<NDArray 5x3 @cpu(0)>
+
+)code" ADD_FILELINE)
+.set_num_inputs(3)
+.set_num_outputs(1)
+.set_attr<nnvm::FInferShape>("FInferShape", IndexCopyShape)
+.set_attr<nnvm::FInferType>("FInferType", IndexCopyType)
+.set_attr<FCompute>("FCompute<cpu>", IndexCopyCompute<cpu>)
+.add_argument("old_tensor", "NDArray", "Old tensor")
+.add_argument("index", "NDArray", "Index vector")
+.add_argument("new_tensor", "NDArray", "New tensor to be copied");
+
+}  // namespace op
+}  // namespace mxnet

src/operator/contrib/index_copy.cu

@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * \file index_copy.cc
+ * \brief
+ */
+#include "./index_copy-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_contrib_index_copy)
+.set_attr<FCompute>("FCompute<gpu>", IndexCopyCompute<gpu>);
+
+}
+}

tests/python/unittest/test_operator.py

@@ -4639,6 +4639,16 @@ def test_quantization_op():
     assert same(qa.asnumpy(), qa_real.asnumpy())
     assert same(a_.asnumpy(),  a_real.asnumpy())
 
+@with_seed()
+def test_index_copy():
+    x = mx.nd.zeros((5,3))
+    t = mx.nd.array([[1,2,3],[4,5,6],[7,8,9]])
+    index = mx.nd.array([0,4,2])
+    out = mx.nd.contrib.index_copy(x, index, t)
+
+    tensor = mx.nd.array([[1,2,3],[0,0,0],[7,8,9],[0,0,0],[4,5,6]])
+
+    assert same(out.asnumpy(), tensor.asnumpy())
 
 @with_seed()
 def test_div_sqrt_dim():