showing error message when user attempts to pass large ndarrays as in…

…puts to operators or large shape of inputs for memory allocation.

python/mxnet/ndarray/ndarray.py

@@ -155,6 +155,11 @@ def _new_alloc_handle(shape, ctx, delay_alloc, dtype=mx_real_t):
             ctypes.c_int(int(_DTYPE_NP_TO_MX[np.dtype(dtype).type])),
             ctypes.byref(hdl)))
     else:
+        size = 1
+        for idx in shape:
+            size = size * idx
+        if size > 2**32:
+            print("Size of tensor you are trying to allocate is larger than 2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1")
         check_call(_LIB.MXNDArrayCreateEx(
             c_array_buf(mx_uint, native_array('I', shape)),
             mx_uint(len(shape)),

src/c_api/c_api_ndarray.cc

@@ -54,7 +54,10 @@ void SetNDInputsOutputs(const nnvm::Op* op,
   ndinputs->clear();
   ndinputs->reserve(num_inputs);
   for (int i = 0; i < num_inputs; ++i) {
-    ndinputs->emplace_back(reinterpret_cast<NDArray*>(inputs[i]));
+    NDArray* inp = reinterpret_cast<NDArray*>(inputs[i]);
+    CHECK_LT(inp->shape().Size(), 1 >> 31) << "Size of tensor you are trying to allocate is larger than "
+                               "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
+    ndinputs->emplace_back(inp);
   }
 
   ndoutputs->clear();
@@ -97,6 +100,7 @@ void MXImperativeInvokeImpl(AtomicSymbolCreator creator,
   imperative::SetNumOutputs(op, attrs, num_inputs, &infered_num_outputs, &num_visible_outputs);
 
   std::vector<NDArray*> ndinputs, ndoutputs;
+
   SetNDInputsOutputs(op, &ndinputs, &ndoutputs, num_inputs, inputs,
       num_outputs, infered_num_outputs, num_visible_outputs, outputs);
 

src/operator/tensor/init_op.h

@@ -336,6 +336,8 @@ inline bool InitStorageType(const nnvm::NodeAttrs& attrs,
 template <bool is_integer = false, typename ValueType, typename xpu>
 void Fill(mshadow::Stream<xpu> *s, const TBlob& b, const OpReqType req, ValueType val) {
   // If b is a zero-size tensor, do nothing.
+  CHECK_LT(b.Size(), 1 >> 31) << "Size of tensor you are trying to allocate is larger than "
+                                 "2^32 elements. Please build with flag USE_INT64_TENSOR_SIZE=1";
   if (b.Size() == 0) return;
   if (req != kNullOp) {
     const size_t size = b.Size();