From 014a13751f2002775bf4b8b813e1d27b1511397f Mon Sep 17 00:00:00 2001
From: Rohit Kumar Srivastava <srivastava.141@buckeyemail.osu.edu>
Date: Mon, 14 Oct 2019 22:00:52 +0000
Subject: [PATCH] adding large tensor support for add_n and tests for more ops

---
 src/operator/tensor/elemwise_sum.h |  8 ++---
 tests/nightly/test_large_array.py  | 52 ++++++++++++++++++++++++++++++
 tests/nightly/test_large_vector.py | 52 ++++++++++++++++++++++++++++++
 3 files changed, 108 insertions(+), 4 deletions(-)

diff --git a/src/operator/tensor/elemwise_sum.h b/src/operator/tensor/elemwise_sum.h
index 08f57d15fbf8..e89e9d799903 100644
--- a/src/operator/tensor/elemwise_sum.h
+++ b/src/operator/tensor/elemwise_sum.h
@@ -39,15 +39,15 @@ namespace op {
 
 struct Sum {
   template<typename DType>
-  MSHADOW_XINLINE static DType sum(int i, const DType* a) {
+  MSHADOW_XINLINE static DType sum(index_t i, const DType* a) {
     return a[i];
   }
   template<typename DType, typename... DTypes>
-  MSHADOW_XINLINE static DType sum(int i, const DType* a, const DTypes... b) {
+  MSHADOW_XINLINE static DType sum(index_t i, const DType* a, const DTypes... b) {
     return a[i] + sum(i, b...);
   }
   template<typename DType, typename... DTypes>
-  MSHADOW_XINLINE static void Map(int i, DType* out, const OpReqType req, const DType* in0,
+  MSHADOW_XINLINE static void Map(index_t i, DType* out, const OpReqType req, const DType* in0,
     const DTypes... ins) {
     KERNEL_ASSIGN(out[i], req, sum(i, in0, ins...));
   }
@@ -64,7 +64,7 @@ void ElementWiseSumCompute_(const nnvm::NodeAttrs& attrs,
   size_t size = in_data.size();
   Stream<xpu> *s = ctx.get_stream<xpu>();
   DType* out_dptr = out_data[0].dptr<DType>();
-  int out_size = static_cast<int>((out_data[0].Size() + DataType<DType>::kLanes - 1)
+  index_t out_size = static_cast<index_t>((out_data[0].Size() + DataType<DType>::kLanes - 1)
                                   /DataType<DType>::kLanes);
   switch (size) {
     case 2: {
diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
index e51e220c232f..8ee633ec5553 100644
--- a/tests/nightly/test_large_array.py
+++ b/tests/nightly/test_large_array.py
@@ -18,6 +18,7 @@
 import math
 import numpy as np
 import mxnet as mx
+import os
 
 from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, default_context, check_symbolic_forward, create_2d_tensor
 from mxnet import gluon, nd
@@ -1199,6 +1200,57 @@ def test_full():
     assert a[-1][-1] == 3
 
 
+def test_load_save():
+    x = create_2d_tensor(SMALL_Y, LARGE_X)
+    nd.save('large_tensor', [x])
+    y = nd.load('large_tensor')
+    y = y[0]
+    assert x[0][0] == y[0][0]
+    assert x[-1][-1]== y[-1][-1]
+    os.remove('large_tensor')
+
+
+def test_add_n():
+    x = [nd.ones(LARGE_X) for j in range(SMALL_Y)]
+    y = nd.add_n(*x)
+    assert y[0] == SMALL_Y
+    assert y[-1] == SMALL_Y
+
+
+def test_modulo():
+    x = mx.nd.ones((SMALL_Y, LARGE_X))*6
+    y = mx.nd.ones(LARGE_X)*4
+    z = (x%y)
+    assert z[0][0] == 2
+    assert z[-1][-1] == 2
+    x = mx.nd.ones((SMALL_Y, LARGE_X))*5
+    z = nd.modulo(x,y)
+    assert z[0][0] == 1
+    assert z[-1][-1] == 1
+
+
+def test_maximum():
+    x = mx.nd.ones((SMALL_Y, LARGE_X))*3
+    y = mx.nd.ones(LARGE_X)*4
+    z = nd.maximum(x, y)
+    assert z[0][0] == 4
+    assert z[-1][-1] == 4
+    z = nd.maximum(x, 5)
+    assert z[0][0] == 5
+    assert z[-1][-1] == 5
+
+
+def test_minimum():
+    x = mx.nd.ones((SMALL_Y, LARGE_X))*3
+    y = mx.nd.ones(LARGE_X)*2
+    z = nd.minimum(x, y)
+    assert z[0][0] == 2
+    assert z[-1][-1] == 2
+    z = nd.minimum(x, 5)
+    assert z[0][0] == 3
+    assert z[-1][-1] == 3
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()
diff --git a/tests/nightly/test_large_vector.py b/tests/nightly/test_large_vector.py
index aa6cb3d75b37..5fe58cd2e49b 100644
--- a/tests/nightly/test_large_vector.py
+++ b/tests/nightly/test_large_vector.py
@@ -18,6 +18,7 @@
 import math
 import numpy as np
 import mxnet as mx
+import os
 
 from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, create_vector
 from mxnet import gluon, nd
@@ -708,6 +709,57 @@ def test_full():
     assert a[-1] == 3
 
 
+def test_load_save():
+    x = create_vector(size=LARGE_X)
+    nd.save('large_vector', [x])
+    y = nd.load('large_vector')
+    y = y[0]
+    assert x[0] == y[0]
+    assert x[-1] == y[-1]
+    os.remove('large_vector')
+
+
+def test_add_n():
+    x = [nd.ones(LARGE_X)]
+    y = nd.add_n(*x)
+    assert y[0] == 1
+    assert y[-1] == 1
+
+
+def test_modulo():
+    x = mx.nd.ones(LARGE_X)*6
+    y = mx.nd.ones(LARGE_X)*4
+    z = (x%y)
+    assert z[0] == 2
+    assert z[-1] == 2
+    x = mx.nd.ones(LARGE_X)*5
+    z = nd.modulo(x,y)
+    assert z[0] == 1
+    assert z[-1] == 1
+
+
+def test_maximum():
+    x = mx.nd.ones(LARGE_X)*3
+    y = mx.nd.ones(LARGE_X)*4
+    z = nd.maximum(x, y)
+    assert z[0] == 4
+    assert z[-1] == 4
+    z = nd.maximum(x, 5)
+    assert z[0] == 5
+    assert z[-1] == 5
+
+
+def test_minimum():
+    x = mx.nd.ones(LARGE_X)*3
+    y = mx.nd.ones(LARGE_X)*2
+    z = nd.minimum(x, y)
+    assert z[0] == 2
+    assert z[-1] == 2
+    z = nd.minimum(x, 5)
+    assert z[0] == 3
+    assert z[-1] == 3
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()