You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by zh...@apache.org on 2020/12/03 01:19:05 UTC

[incubator-mxnet] branch master updated: Numpy unique repeat indices large tensor checks (#19382)

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

zha0q1 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 619eab2  Numpy unique repeat indices large tensor checks (#19382)
619eab2 is described below

commit 619eab2657d10884a3db6d28d3bc10e14fd3534c
Author: Zhaoqi Zhu <zh...@gmail.com>
AuthorDate: Wed Dec 2 17:18:06 2020 -0800

    Numpy unique repeat indices large tensor checks (#19382)
    
    * add size checks to repeat and unique
    
    * add size checks
    
    * Update test_np_large_array.py
    
    * Update np_repeat_op-inl.h
    
    * Update np_init_op.cc
    
    * Update np_init_op.cc
    
    * Update np_repeat_op-inl.h
    
    * Update np_unique_op.cc
---
 src/operator/numpy/np_init_op.cc      |  4 +++-
 src/operator/numpy/np_repeat_op-inl.h |  4 ++++
 src/operator/numpy/np_unique_op.cc    |  2 ++
 tests/nightly/test_np_large_array.py  | 18 +++++++++++++++++-
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/src/operator/numpy/np_init_op.cc b/src/operator/numpy/np_init_op.cc
index e30e977..6e7aca2 100644
--- a/src/operator/numpy/np_init_op.cc
+++ b/src/operator/numpy/np_init_op.cc
@@ -47,11 +47,13 @@ inline bool NumpyIndicesShape(const nnvm::NodeAttrs& attrs,
     << "_npi_indices dimensions the number of dim must not be less than 0";
   mxnet::TShape param_dim = param.dimensions;
   if (!shape_is_known(param_dim)) return false;
+  CHECK_LT(param_dim.Size(), INT32_MAX) << "ValueError: np.indices does not support large"
+     << " input tensors (containing >= 2^31 elements).";
   const int indim = param.dimensions.ndim();
   mxnet::TShape ret(indim + 1, -1);
   ret[0] = indim;
   for (int i = 1; i < indim + 1; ++i) {
-    ret[i] = param.dimensions[i-1];
+    ret[i] = param_dim[i-1];
   }
   SHAPE_ASSIGN_CHECK(*out_shapes, 0, ret);
   return shape_is_known(out_shapes->at(0));
diff --git a/src/operator/numpy/np_repeat_op-inl.h b/src/operator/numpy/np_repeat_op-inl.h
index 638f1de..aa51d08 100644
--- a/src/operator/numpy/np_repeat_op-inl.h
+++ b/src/operator/numpy/np_repeat_op-inl.h
@@ -110,10 +110,14 @@ inline bool RepeatsOpShape(const nnvm::NodeAttrs& attrs,
         shape[i] = ishape[i];
       }
     }
+    CHECK_LT(shape.Size(), INT32_MAX) << "ValueError: np.repeat does not support large"
+      << " input tensors (containing >= 2^31 elements).";
     SHAPE_ASSIGN_CHECK(*out_attrs, 0, shape);
   } else {  // If axis is not input by user, return a flat 1D array of size = repeats
     repeats = param.repeats.value().ndim() == 1 ? ishape.Size() * repeats : repeats;
     mxnet::TShape shape(1, repeats);
+    CHECK_LT(shape.Size(), INT32_MAX) << "ValueError: np.repeat does not support large"
+      << " input tensors (containing >= 2^31 elements).";
     SHAPE_ASSIGN_CHECK(*out_attrs, 0, shape);
   }
   return shape_is_known(out_attrs->at(0));
diff --git a/src/operator/numpy/np_unique_op.cc b/src/operator/numpy/np_unique_op.cc
index 39a84ba..7076b44 100644
--- a/src/operator/numpy/np_unique_op.cc
+++ b/src/operator/numpy/np_unique_op.cc
@@ -348,6 +348,8 @@ void NumpyUniqueCPUForward(const nnvm::NodeAttrs& attrs,
       const_cast<NDArray &>(outputs[output_flag]).Init(shape_0);
     }
   } else {
+    CHECK_LT(inputs[0].shape().Size(), INT32_MAX) << "ValueError: np.unique does not support large"
+      << " input tensors (containing >= 2^31).";
     if (!param.axis.has_value()) {
       NumpyUniqueCPUNoneAxisImpl(param, ctx, inputs, req, outputs);
     } else {
diff --git a/tests/nightly/test_np_large_array.py b/tests/nightly/test_np_large_array.py
index a1a34d8..dd4a94c 100644
--- a/tests/nightly/test_np_large_array.py
+++ b/tests/nightly/test_np_large_array.py
@@ -2132,6 +2132,23 @@ def test_dsplit():
 
 
 @use_np
+def test_unique():
+    inp = np.zeros((2, HALF_INT_OVERFLOW))
+    assertRaises(ValueError, np.unique, inp, axis=1)
+
+
+@use_np
+def test_repeat():
+    inp = np.ones((2, HALF_INT_OVERFLOW))
+    assertRaises(ValueError, np.repeat, inp, repeats=2, axis=1)
+
+
+@use_np
+def test_indices():
+    assertRaises(ValueError, np.indices, (2, HALF_INT_OVERFLOW))
+
+
+@use_np    
 def test_tril_indices():
     N = 2**16
     data = np.tril_indices(N, -1)
@@ -2328,4 +2345,3 @@ def test_insert():
     assert out[0, 1] == 1 and out[-1, 1] == 2
     assert out2[1] == 5 and out2[2] == 6
     assertRaises(MXNetError, np.insert, arr=inp3, obj=np.array([2, 2], dtype=np.int64), values=np.array([5, 6]))
-