[GitHub] [incubator-mxnet] Zha0q1 commented on a change in pull request #19357: [TEST] Added large tensor tests to verify support for split ops

[GitBox <gi...@apache.org>]

Zha0q1 commented on a change in pull request #19357:
URL: https://github.com/apache/incubator-mxnet/pull/19357#discussion_r505927937



##########
File path: tests/nightly/test_np_large_array.py
##########
@@ -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)

Review comment:
       use create_2d_np_tensor here?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org