From 99e69e68258d832f9f866988c1f25e0a579ddd69 Mon Sep 17 00:00:00 2001
From: Rohit Kumar Srivastava <srivastava.141@osu.edu>
Date: Sun, 2 Jun 2019 15:06:33 -0700
Subject: [PATCH] [MXNET-1410]Adding Large Tensor Support for tensor transpose
 (#15059)

---
 src/operator/tensor/matrix_op-inl.h |  8 ++++----
 tests/nightly/test_large_array.py   | 27 +++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/src/operator/tensor/matrix_op-inl.h b/src/operator/tensor/matrix_op-inl.h
index 50cb1ae4d1df..5cd7bf6652d3 100644
--- a/src/operator/tensor/matrix_op-inl.h
+++ b/src/operator/tensor/matrix_op-inl.h
@@ -1950,10 +1950,10 @@ struct ReverseParam : public dmlc::Parameter<ReverseParam> {
 #define REVERSE_MAX_DIM 10U
 
 struct reverse {
-  MSHADOW_XINLINE static int ReverseIndex(index_t idx,
-                                          index_t nreversedim,
-                                          const index_t * stride_,
-                                          const index_t * trailing_) {
+  MSHADOW_XINLINE static index_t ReverseIndex(index_t idx,
+                                              index_t nreversedim,
+                                              const index_t * stride_,
+                                              const index_t * trailing_) {
     index_t outputIndex = idx;
     for (index_t i = 0; i < nreversedim; ++i) {
       const index_t low = outputIndex % trailing_[i];
diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
index 286f2e473bdd..8c80e9e22488 100644
--- a/tests/nightly/test_large_array.py
+++ b/tests/nightly/test_large_array.py
@@ -292,6 +292,33 @@ def test_unravel_index():
     assert (indices_2d.asnumpy() == np.array(original_2d_indices)).all()
 
 
+def create_2d_tensor(rows, columns):
+    a = np.arange(0, rows).reshape(rows, 1)
+    b = np.broadcast_to(a, shape=(a.shape[0], columns))
+    return nd.array(b, dtype=np.int64)
+
+
+def test_transpose():
+    b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y)
+    t = b.T
+    assert t.shape == (SMALL_Y, LARGE_X)
+    assert np.sum(t[:, -1].asnumpy() == (LARGE_X - 1)) == b.shape[1]
+
+
+def test_swapaxes():
+    b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y)
+    t = nd.swapaxes(b, dim1=0, dim2=1)
+    assert t.shape == (SMALL_Y, LARGE_X)
+    assert np.sum(t[:, -1].asnumpy() == (LARGE_X - 1)) == b.shape[1]
+
+
+def test_flip():
+    b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y)
+    t = nd.flip(b, axis=0)
+    assert t.shape == (LARGE_X, SMALL_Y)
+    assert np.sum(t[-1, :].asnumpy() == 0) == b.shape[1]
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()