diff --git a/src/operator/contrib/index_copy-inl.h b/src/operator/contrib/index_copy-inl.h index b97138a88f97..923fb0f4f138 100644 --- a/src/operator/contrib/index_copy-inl.h +++ b/src/operator/contrib/index_copy-inl.h @@ -32,6 +32,7 @@ #include "../elemwise_op_common.h" #include "../mshadow_op.h" #include "../mxnet_op.h" +#include "../tensor/init_op.h" namespace mxnet { namespace op { @@ -83,12 +84,12 @@ void IndexCopyForward(const nnvm::NodeAttrs& attrs, }); } -template struct index_copy_backward { template MSHADOW_XINLINE static void Map(int i, int dim, int index_size, + int req1, int req2, DType* out_grad, IType* index, DType* in_grad_1, @@ -98,12 +99,12 @@ struct index_copy_backward { int idx = static_cast(index[p]); if (i >= idx*dim && i < (idx+1)*dim) { int offset = i - idx*dim; - KERNEL_ASSIGN(in_grad_2[p*dim+offset], req, out_grad[i]); + KERNEL_ASSIGN(in_grad_2[p*dim+offset], req2, out_grad[i]); return; } } // Copy to in_grad_1 - KERNEL_ASSIGN(in_grad_1[i], req, out_grad[i]); + KERNEL_ASSIGN(in_grad_1[i], req1, out_grad[i]); } }; @@ -122,18 +123,19 @@ void IndexCopyBackward(const nnvm::NodeAttrs& attrs, const TBlob& in_grad_2 = outputs[2]; int dim = inputs[3].Size() / inputs[2].Size(); int index_size = inputs[2].Size(); + Fill(s, outputs[0], req[0], 0); + Fill(s, outputs[2], req[2], 0); // index_copy_backward MSHADOW_TYPE_SWITCH(out_grad.type_flag_, DType, { MSHADOW_TYPE_SWITCH(index.type_flag_, IType, { - MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, { - mxnet_op::Kernel, xpu>::Launch(s, + mxnet_op::Kernel::Launch(s, out_grad.Size(), dim, index_size, + req[0], req[2], out_grad.dptr(), index.dptr(), in_grad_1.dptr(), in_grad_2.dptr()); - }); }); }); } diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 80a83df45da0..c59cdfb53896 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -4768,24 +4768,25 @@ 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], dtype=np.int64) + tensor = mx.nd.array([[1,2,3],[0,0,0],[7,8,9],[0,0,0],[4,5,6]]) + x_grad = mx.nd.array([[0,0,0],[1,1,1],[0,0,0],[1,1,1],[0,0,0]]) + t_grad = mx.nd.array([[1,1,1],[1,1,1],[1,1,1]]) - x.attach_grad() t.attach_grad() - index.attach_grad() - with mx.autograd.record(): out = mx.nd.contrib.index_copy(x, index, t) out.backward() + assert same(out.asnumpy(), tensor.asnumpy()) + assert same(t.grad.asnumpy(), t_grad.asnumpy()) - tensor = mx.nd.array([[1,2,3],[0,0,0],[7,8,9],[0,0,0],[4,5,6]]) - x_grad = mx.nd.array([[0,0,0],[1,1,1],[0,0,0],[1,1,1],[0,0,0]]) - t_grad = mx.nd.array([[1,1,1],[1,1,1],[1,1,1]]) - index_grad = mx.nd.array([0,0,0]) - + x.attach_grad() + t.attach_grad() + with mx.autograd.record(): + out = mx.nd.contrib.index_copy(x, index, t) + out.backward() assert same(out.asnumpy(), tensor.asnumpy()) assert same(x.grad.asnumpy(), x_grad.asnumpy()) assert same(t.grad.asnumpy(), t_grad.asnumpy()) - assert same(index.grad.asnumpy(), index_grad.asnumpy()) @with_seed() def test_div_sqrt_dim():