[incubator-mxnet] branch master updated: [FEATURE]Enable large tensor support for batch flatten (#19063)

[sa...@apache.org]

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

samskalicky 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 f732530  [FEATURE]Enable large tensor support for batch flatten (#19063)
f732530 is described below

commit f732530c8e1f8dcc11134e935430e3793ddbf4c8
Author: Rohit Kumar Srivastava <sr...@osu.edu>
AuthorDate: Wed Sep 2 18:44:05 2020 -0700

    [FEATURE]Enable large tensor support for batch flatten (#19063)
    
    * Enable large tensor support for batch flatten
    
    * Re-Trigger build
    
    Co-authored-by: Rohit Kumar Srivastava <sr...@buckeyemail.osu.edu>
---
 src/operator/tensor/matrix_op-inl.h  | 2 +-
 tests/nightly/test_np_large_array.py | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/operator/tensor/matrix_op-inl.h b/src/operator/tensor/matrix_op-inl.h
index 209cc36..6e73642 100644
--- a/src/operator/tensor/matrix_op-inl.h
+++ b/src/operator/tensor/matrix_op-inl.h
@@ -243,7 +243,7 @@ inline bool FlattenShape(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(out_attrs->size(), 1U);
   const mxnet::TShape &dshape = (*in_attrs)[0];
   if (!shape_is_known(dshape)) return false;
-  int target_dim = 1;
+  size_t target_dim = 1;
   for (int i = 1; i < dshape.ndim(); ++i) {
     target_dim *= dshape[i];
   }
diff --git a/tests/nightly/test_np_large_array.py b/tests/nightly/test_np_large_array.py
index d2b653b..606b4b5 100644
--- a/tests/nightly/test_np_large_array.py
+++ b/tests/nightly/test_np_large_array.py
@@ -651,7 +651,6 @@ def test_constraint_check():
 
 # broken
 @use_np
-@pytest.mark.skip(reason='Does not support large tensor; to be fixed')
 def test_batch_flatten():
     A = np.ones((2, 1, INT_OVERFLOW))
     A.attach_grad()