Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit 9c2810e

Browse files
aksnzhyeric-haibin-lin
aksnzhy
authored and
eric-haibin-lin
committed
Add index_copy() operator (#12810)
* add index_copy operator * add index_copy op * update index_copy op * add unittest for index_copy() * update index_copy * update index_copy * use mxnet_op::copy * update index_copy * update index_copy * update index_copy * update index_copy test * update index_copy test
1 parent daada21 commit 9c2810e

4 files changed

Lines changed: 305 additions & 0 deletions

File tree

‎src/operator/contrib/index_copy-inl.h‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
/*!
21+
* \file index_copy-inl.h
22+
* \brief implementation of index_copy tensor operation
23+
*/
24+
25+
#ifndef MXNET_OPERATOR_CONTRIB_INDEX_COPY_INL_H_
26+
#define MXNET_OPERATOR_CONTRIB_INDEX_COPY_INL_H_
27+
28+
#include <mxnet/operator_util.h>
29+
#include <vector>
30+
#include <limits>
31+
#include <algorithm>
32+
#include "../elemwise_op_common.h"
33+
#include "../mshadow_op.h"
34+
#include "../mxnet_op.h"
35+
36+
namespace mxnet {
37+
namespace op {
38+
39+
template<int req>
40+
struct index_copy_forward {
41+
template<typename DType, typename IType>
42+
MSHADOW_XINLINE static void Map(int i,
43+
int dim,
44+
IType* index,
45+
DType* new_tensor,
46+
DType* out_tensor) {
47+
DType* out_ptr = out_tensor + static_cast<int>(index[i]) * dim;
48+
DType* new_ptr = new_tensor + i * dim;
49+
for (int idx = 0; idx < dim; ++idx) {
50+
KERNEL_ASSIGN(out_ptr[idx], req, new_ptr[idx]);
51+
}
52+
}
53+
};
54+
55+
template<typename xpu>
56+
void IndexCopyForward(const nnvm::NodeAttrs& attrs,
57+
const OpContext& ctx,
58+
const std::vector<TBlob>& inputs,
59+
const std::vector<OpReqType>& req,
60+
const std::vector<TBlob>& outputs) {
61+
CHECK_EQ(inputs.size(), 3U);
62+
CHECK_EQ(outputs.size(), 1U);
63+
CHECK_EQ(req.size(), 1U);
64+
mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
65+
const TBlob& out = outputs[0];
66+
const TBlob& original_tensor = inputs[0];
67+
const TBlob& idx_vector = inputs[1];
68+
const TBlob& copied_tensor = inputs[2];
69+
int dim = inputs[2].Size() / inputs[1].Size();
70+
// copy original tensor to output
71+
mxnet_op::copy(s, out, original_tensor);
72+
// index copy
73+
MSHADOW_TYPE_SWITCH(out.type_flag_, DType, {
74+
MSHADOW_TYPE_SWITCH(idx_vector.type_flag_, IType, {
75+
MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
76+
mxnet_op::Kernel<index_copy_forward<req_type>, xpu>::Launch(s,
77+
idx_vector.Size(), dim,
78+
idx_vector.dptr<IType>(),
79+
copied_tensor.dptr<DType>(),
80+
out.dptr<DType>());
81+
});
82+
});
83+
});
84+
}
85+
86+
template<int req>
87+
struct index_copy_backward {
88+
template<typename DType, typename IType>
89+
MSHADOW_XINLINE static void Map(int i,
90+
int dim,
91+
int index_size,
92+
DType* out_grad,
93+
IType* index,
94+
DType* in_grad_1,
95+
DType* in_grad_2) {
96+
// Copy to in_grad_2
97+
for (int p = 0; p < index_size; ++p) {
98+
int idx = static_cast<int>(index[p]);
99+
if (i >= idx*dim && i < (idx+1)*dim) {
100+
int offset = i - idx*dim;
101+
KERNEL_ASSIGN(in_grad_2[p*dim+offset], req, out_grad[i]);
102+
return;
103+
}
104+
}
105+
// Copy to in_grad_1
106+
KERNEL_ASSIGN(in_grad_1[i], req, out_grad[i]);
107+
}
108+
};
109+
110+
template<typename xpu>
111+
void IndexCopyBackward(const nnvm::NodeAttrs& attrs,
112+
const OpContext& ctx,
113+
const std::vector<TBlob>& inputs,
114+
const std::vector<OpReqType>& req,
115+
const std::vector<TBlob>& outputs) {
116+
CHECK_EQ(inputs.size(), 4U);
117+
CHECK_EQ(outputs.size(), 3U);
118+
mshadow::Stream<xpu> *s = ctx.get_stream<xpu>();
119+
const TBlob& out_grad = inputs[0];
120+
const TBlob& index = inputs[2];
121+
const TBlob& in_grad_1 = outputs[0];
122+
const TBlob& in_grad_2 = outputs[2];
123+
int dim = inputs[3].Size() / inputs[2].Size();
124+
int index_size = inputs[2].Size();
125+
// index_copy_backward
126+
MSHADOW_TYPE_SWITCH(out_grad.type_flag_, DType, {
127+
MSHADOW_TYPE_SWITCH(index.type_flag_, IType, {
128+
MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
129+
mxnet_op::Kernel<index_copy_backward<req_type>, xpu>::Launch(s,
130+
out_grad.Size(),
131+
dim, index_size,
132+
out_grad.dptr<DType>(),
133+
index.dptr<IType>(),
134+
in_grad_1.dptr<DType>(),
135+
in_grad_2.dptr<DType>());
136+
});
137+
});
138+
});
139+
}
140+
141+
inline bool IndexCopyShape(const nnvm::NodeAttrs& attrs,
142+
std::vector<TShape> *in_attrs,
143+
std::vector<TShape> *out_attrs) {
144+
// inputs[0]: original tensor
145+
// inputs[1]: index vector
146+
// inputs[2]: copied tensor
147+
CHECK_EQ(in_attrs->size(), 3U);
148+
// outputs[0]: a new tensor
149+
CHECK_EQ(out_attrs->size(), 1U);
150+
// inputs[1] must be a vector
151+
CHECK_EQ(in_attrs->at(1).ndim(), 1);
152+
// Shape matching
153+
CHECK_EQ(in_attrs->at(0).ndim(), in_attrs->at(2).ndim());
154+
for (size_t i = 0; i < in_attrs->at(0).ndim(); ++i) {
155+
if (i == 0) {
156+
CHECK_GE(in_attrs->at(0)[i], in_attrs->at(2)[i]);
157+
} else {
158+
CHECK_EQ(in_attrs->at(0)[i], in_attrs->at(2)[i]);
159+
}
160+
}
161+
// The the length of the fitrst dim of copied tensor
162+
// must equal to the size of index vector
163+
CHECK_EQ(in_attrs->at(1)[0], in_attrs->at(2)[0]);
164+
SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
165+
SHAPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
166+
return out_attrs->at(0).ndim() != 0U &&
167+
out_attrs->at(0).Size() != 0U;
168+
}
169+
170+
} // namespace op
171+
} // namespace mxnet
172+
173+
#endif // MXNET_OPERATOR_CONTRIB_INDEX_COPY_INL_H_

‎src/operator/contrib/index_copy.cc‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
/*!
21+
* \file index_copy.cc
22+
* \brief
23+
*/
24+
#include "./index_copy-inl.h"
25+
26+
namespace mxnet {
27+
namespace op {
28+
29+
NNVM_REGISTER_OP(_contrib_index_copy)
30+
.describe(R"code(Copies the elements of a `new_tensor` into the `old_tensor` by
31+
selecting the indices in the order given in `index`. The output will be a new tensor
32+
contains the rest elements of old tensor and the copied elements of new tensor.
33+
For example, if `index[i] == j`, then the `i`th row of `new_tensor` is copied to the
34+
`j`th row of output.
35+
36+
The `index` must be a vector and it must have the same size with the `0`th dimimention of
37+
`new_tensor`. Also, the `0`th dimimention of old_tensor must `>=` the `0`th dimimention of
38+
`new_tensor`, or an error will be raised.
39+
40+
Examples::
41+
42+
x = mx.nd.zeros((5,3))
43+
t = mx.nd.array([[1,2,3],[4,5,6],[7,8,9]])
44+
index = mx.nd.array([0,4,2])
45+
46+
mx.nd.contrib.index_copy(x, index, t)
47+
48+
[[1. 2. 3.]
49+
[0. 0. 0.]
50+
[7. 8. 9.]
51+
[0. 0. 0.]
52+
[4. 5. 6.]]
53+
<NDArray 5x3 @cpu(0)>
54+
55+
)code" ADD_FILELINE)
56+
.set_num_inputs(3)
57+
.set_num_outputs(1)
58+
.set_attr<nnvm::FInferShape>("FInferShape", IndexCopyShape)
59+
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<3, 1>)
60+
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_contrib_backward_index_copy"})
61+
.set_attr<FCompute>("FCompute<cpu>", IndexCopyForward<cpu>)
62+
.add_argument("old_tensor", "NDArray-or-Symbol", "Old tensor")
63+
.add_argument("index_vector", "NDArray-or-Symbol", "Index vector")
64+
.add_argument("new_tensor", "NDArray-or-Symbol", "New tensor to be copied");
65+
66+
NNVM_REGISTER_OP(_contrib_backward_index_copy)
67+
.set_num_inputs(4)
68+
.set_num_outputs(3)
69+
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
70+
.set_attr<FCompute>("FCompute<cpu>", IndexCopyBackward<cpu>);
71+
72+
} // namespace op
73+
} // namespace mxnet

‎src/operator/contrib/index_copy.cu‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
/*!
21+
* \file index_copy.cc
22+
* \brief
23+
*/
24+
#include "./index_copy-inl.h"
25+
26+
namespace mxnet {
27+
namespace op {
28+
29+
NNVM_REGISTER_OP(_contrib_index_copy)
30+
.set_attr<FCompute>("FCompute<gpu>", IndexCopyForward<gpu>);
31+
32+
NNVM_REGISTER_OP(_contrib_backward_index_copy)
33+
.set_attr<FCompute>("FCompute<gpu>", IndexCopyBackward<gpu>);
34+
35+
} // namespace op
36+
} // namespace mxnet

‎tests/python/unittest/test_operator.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4762,6 +4762,29 @@ def test_quantization_op():
47624762
assert same(qa.asnumpy(), qa_real.asnumpy())
47634763
assert same(a_.asnumpy(), a_real.asnumpy())
47644764

4765+
@with_seed()
4766+
def test_index_copy():
4767+
x = mx.nd.zeros((5,3))
4768+
t = mx.nd.array([[1,2,3],[4,5,6],[7,8,9]])
4769+
index = mx.nd.array([0,4,2])
4770+
4771+
x.attach_grad()
4772+
t.attach_grad()
4773+
index.attach_grad()
4774+
4775+
with mx.autograd.record():
4776+
out = mx.nd.contrib.index_copy(x, index, t)
4777+
out.backward()
4778+
4779+
tensor = mx.nd.array([[1,2,3],[0,0,0],[7,8,9],[0,0,0],[4,5,6]])
4780+
x_grad = mx.nd.array([[0,0,0],[1,1,1],[0,0,0],[1,1,1],[0,0,0]])
4781+
t_grad = mx.nd.array([[1,1,1],[1,1,1],[1,1,1]])
4782+
index_grad = mx.nd.array([0,0,0])
4783+
4784+
assert same(out.asnumpy(), tensor.asnumpy())
4785+
assert same(x.grad.asnumpy(), x_grad.asnumpy())
4786+
assert same(t.grad.asnumpy(), t_grad.asnumpy())
4787+
assert same(index.grad.asnumpy(), index_grad.asnumpy())
47654788

47664789
@with_seed()
47674790
def test_div_sqrt_dim():

0 commit comments

Comments
 (0)