[incubator-mxnet] branch master updated: adding large tensor tests to verify support for split, hsplit, vsplit, dsplit (#19357)

[sk...@apache.org]

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

skm 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 7bbe928  adding large tensor tests to verify support for split, hsplit, vsplit, dsplit (#19357)
7bbe928 is described below

commit 7bbe92834e47672d8136b0d5b384978df55f0fc7
Author: Rohit Kumar Srivastava <sr...@osu.edu>
AuthorDate: Fri Oct 16 12:47:40 2020 -0700

    adding large tensor tests to verify support for split, hsplit, vsplit, dsplit (#19357)
    
    Co-authored-by: Rohit Kumar Srivastava <sr...@buckeyemail.osu.edu>
---
 python/mxnet/test_utils.py           |  5 +++
 tests/nightly/test_np_large_array.py | 65 +++++++++++++++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/python/mxnet/test_utils.py b/python/mxnet/test_utils.py
index 0b9e181..01bdaad 100644
--- a/python/mxnet/test_utils.py
+++ b/python/mxnet/test_utils.py
@@ -377,6 +377,11 @@ def assign_each2(input1, input2, function):
 
     return output
 
+def create_2d_np_tensor(rows, columns, dtype=np.int64):
+    inp = mx.np.arange(0, rows, dtype=dtype).reshape(rows, 1)
+    inp = mx.np.broadcast_to(inp, shape=(inp.shape[0], columns))
+    return inp
+
 # For testing Large Tensors having total size > 2^32 elements
 def create_2d_tensor(rows, columns, dtype=np.int64):
     a = mx.nd.arange(0, rows, dtype=dtype).reshape(rows, 1)
diff --git a/tests/nightly/test_np_large_array.py b/tests/nightly/test_np_large_array.py
index c158977..61403b0 100644
--- a/tests/nightly/test_np_large_array.py
+++ b/tests/nightly/test_np_large_array.py
@@ -25,7 +25,7 @@ import mxnet as mx
 curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
 sys.path.append(os.path.join(curr_path, '../python/unittest/'))
 
-from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, default_context, check_symbolic_forward, create_2d_tensor, use_np
+from mxnet.test_utils import rand_ndarray, assert_almost_equal, rand_coord_2d, default_context, check_symbolic_forward, create_2d_np_tensor, use_np
 from mxnet import gluon, np, npx
 from common import with_seed
 import pytest
@@ -2001,3 +2001,66 @@ def test_vstack():
     assert out2[0, -1] == 0 and out2[1, -1] == 1
     assert inp2.grad.shape == inp2.shape
     assert inp2.grad[-1, -1] == 1
+
+
+@use_np
+def test_split():
+    inp = np.ones((INT_OVERFLOW, 2))
+    inp[INT_OVERFLOW // 2] = 2
+    inp.attach_grad()
+    with mx.autograd.record():
+        out = np.split(inp, 2, axis = 0)
+        out[1].backward()
+    assert out[0].shape == (INT_OVERFLOW // 2, 2)
+    assert out[1].shape == (INT_OVERFLOW // 2, 2)
+    assert out[0][0, 0] == 1
+    assert out[1][0, 0] == 2
+    assert inp.grad.shape == inp.shape
+    assert inp.grad[0][0] == 0 and inp.grad[-1][-1] == 1
+
+
+@use_np
+def test_hsplit():
+    inp = create_2d_np_tensor(rows=INT_OVERFLOW, columns=4)
+    inp.attach_grad()
+    with mx.autograd.record():
+        out = np.hsplit(inp, 2)
+        out[1].backward()
+    assert out[1].shape == (INT_OVERFLOW, 2)
+    assert out[0].shape == (INT_OVERFLOW, 2)
+    assert out[1][-1][0] == INT_OVERFLOW-1
+    assert out[0][-1][1] == INT_OVERFLOW-1
+    assert inp.grad.shape == inp.shape
+    assert inp.grad[0][0] == 0 and inp.grad[-1][-1] == 1
+
+
+@use_np
+def test_vsplit():
+    inp = create_2d_np_tensor(rows=INT_OVERFLOW, columns=4)
+    inp.attach_grad()
+    with mx.autograd.record():
+        out = np.vsplit(inp, 2)
+        out[1].backward()
+    assert out[1].shape == (INT_OVERFLOW//2, 4)
+    assert out[0].shape == (INT_OVERFLOW//2, 4)
+    assert out[0][-1][0] == INT_OVERFLOW // 2 -1
+    assert out[1][-1][1] == INT_OVERFLOW - 1
+    assert inp.grad.shape == inp.shape
+    assert inp.grad[INT_OVERFLOW//2 - 1][-1] == 0 and inp.grad[-1][-1] == 1
+
+
+@use_np
+def test_dsplit():
+    inp = np.arange(INT_OVERFLOW, dtype=np.int64).reshape(INT_OVERFLOW, 1, 1)
+    inp = np.broadcast_to(inp, shape=(inp.shape[0], 2, 2))
+    inp.attach_grad()
+    with mx.autograd.record():
+        out = np.dsplit(inp, 2)
+        out[1].backward()
+    assert out[1].shape == (INT_OVERFLOW, 2, 1)
+    assert out[0].shape == (INT_OVERFLOW, 2, 1)
+    assert out[0][-1][0][0] == INT_OVERFLOW - 1
+    assert out[1][0][1][0] == 0
+    assert inp.grad.shape == inp.shape
+    assert inp.grad[-1][-1][0] == 0 and inp.grad[0][1][1] == 1
+