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
apache · Dec 3, 2020 · 619eab2 · 619eab2
1 parent e6a855c
commit 619eab2
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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]))
-