[incubator-mxnet] branch v1.x updated: Optimize AddTakeGrad Tensor Sum (#17906) (#18045)

[pa...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

patriczhao pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/v1.x by this push:
     new 3f920ae  Optimize AddTakeGrad Tensor Sum (#17906) (#18045)
3f920ae is described below

commit 3f920ae22c03ba3ff8709c5e695805aca64f2e3a
Author: YixinBao <yi...@intel.com>
AuthorDate: Wed Apr 15 14:15:08 2020 +0800

    Optimize AddTakeGrad Tensor Sum (#17906) (#18045)
---
 3rdparty/mshadow/mshadow/tensor_cpu-inl.h | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/3rdparty/mshadow/mshadow/tensor_cpu-inl.h b/3rdparty/mshadow/mshadow/tensor_cpu-inl.h
index b7ae77f..7531dd9 100755
--- a/3rdparty/mshadow/mshadow/tensor_cpu-inl.h
+++ b/3rdparty/mshadow/mshadow/tensor_cpu-inl.h
@@ -498,9 +498,10 @@ template<bool clip, typename IndexType, typename DType>
 inline void AddTakeGrad(Tensor<cpu, 2, DType> dst,
                         const Tensor<cpu, 1, IndexType>& index,
                         const Tensor<cpu, 2, DType> &src) {
-  const int K = dst.shape_[0];
+  const index_t K = dst.shape_[0];
+  const index_t C = dst.shape_[1];
   for (index_t y = 0; y < index.size(0); ++y) {
-    int j = index[y];
+    index_t j = index[y];
     if (clip) {
       if (j <= 0) j = 0;
       else if (j >= K) j = K - 1;
@@ -508,7 +509,9 @@ inline void AddTakeGrad(Tensor<cpu, 2, DType> dst,
       j %= K;
       if (j < 0) j += K;
     }
-    dst[j] += src[y];
+    for (index_t i = 0; i < C; ++i) {
+      dst[j][i] += src[y][i];
+    }
   }
 }