[incubator-mxnet] branch master updated: Fix nd.pick large array issue (#14082)

[ha...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new ed0b791  Fix nd.pick large array issue (#14082)
ed0b791 is described below

commit ed0b791ed5c2c2dde8b7136302abbbd3a85e89da
Author: Chaitanya Prakash Bapat <ch...@gmail.com>
AuthorDate: Sat Feb 16 00:01:30 2019 -0500

    Fix nd.pick large array issue (#14082)
    
    * large op support
    
    * replaced size_t with index_t for M, added test case
    
    * changed shape
---
 src/operator/tensor/broadcast_reduce_op.h | 12 ++++++------
 tests/nightly/test_large_array.py         |  5 +++++
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/operator/tensor/broadcast_reduce_op.h b/src/operator/tensor/broadcast_reduce_op.h
index 1edcb5a..6aeeadf 100644
--- a/src/operator/tensor/broadcast_reduce_op.h
+++ b/src/operator/tensor/broadcast_reduce_op.h
@@ -1172,12 +1172,12 @@ void L2NormComputeEx(const nnvm::NodeAttrs& attrs,
 template<int ndim, bool clip = true>
 struct pick {
   template<typename DType, typename IType>
-  MSHADOW_XINLINE static void Map(int i, DType* out, const DType* a,
-                                  const IType *idx, int M, int stride,
+  MSHADOW_XINLINE static void Map(index_t i, DType* out, const DType* a,
+                                  const IType *idx, index_t M, int stride,
                                   mshadow::Shape<ndim> bshape,
                                   mshadow::Shape<ndim> sshape) {
     using namespace broadcast;
-    int j = static_cast<int>(idx[i]);
+    index_t j = static_cast<index_t>(idx[i]);
     if (clip) {
       if (j <= 0) j = 0;
       else if (j >= M) j = M - 1;
@@ -1194,12 +1194,12 @@ struct pick {
 template<int ndim, bool clip = true>
 struct pick_grad {
   template<typename DType, typename IType>
-  MSHADOW_XINLINE static void Map(int i, DType* igrad, const DType* ograd,
-                                  const IType *idx, int M, int stride,
+  MSHADOW_XINLINE static void Map(index_t i, DType* igrad, const DType* ograd,
+                                  const IType *idx, index_t M, int stride,
                                   mshadow::Shape<ndim> bshape,
                                   mshadow::Shape<ndim> sshape) {
     using namespace broadcast;
-    int j = static_cast<int>(idx[i]);
+    index_t j = static_cast<index_t>(idx[i]);
     if (clip) {
       if (j <= 0) j = 0;
       else if (j >= M) j = M - 1;
diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
index 696fdb1..0249f44 100644
--- a/tests/nightly/test_large_array.py
+++ b/tests/nightly/test_large_array.py
@@ -145,6 +145,11 @@ def test_where():
     res = nd.sparse.where(csr_cond, a, b)
     assert np.sum(res[0].asnumpy() == 1) == b.shape[1]
 
+def test_pick():
+    a = mx.nd.ones(shape=(256*35, 1024*1024))
+    b = mx.nd.ones(shape=(256*35,))
+    res = mx.nd.pick(a,b)
+    assert res.shape == b.shape
 
 if __name__ == '__main__':
     import nose