[MXNET-1410]Adding Large Tensor Support for tensor transpose

apache · May 24, 2019 · cc20e34 · cc20e34
1 parent 66aa983
commit cc20e34
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/src/operator/tensor/matrix_op-inl.h b/src/operator/tensor/matrix_op-inl.h
@@ -1950,7 +1950,7 @@ struct ReverseParam : public dmlc::Parameter<ReverseParam> {
 #define REVERSE_MAX_DIM 10U
 
 struct reverse {
-  MSHADOW_XINLINE static int ReverseIndex(index_t idx,
+  MSHADOW_XINLINE static index_t ReverseIndex(index_t idx,
                                           index_t nreversedim,
                                           const index_t * stride_,
                                           const index_t * trailing_) {

diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py
@@ -279,6 +279,19 @@ def test_diag():
     assert_almost_equal(r.asnumpy(), np.diag(a_np, k=k))
 
 
+def test_transpose():
+    a = nd.arange(0, LARGE_X).reshape(LARGE_X, 1)
+    b = nd.broadcast_to(a, shape=(a.shape[0], SMALL_Y))
+    t = b.T
+    assert t.shape == (SMALL_Y, LARGE_X)
+    assert np.sum(t[:,-1].asnumpy() == LARGE_X) == b.shape[1]
+    t = nd.swapaxes(b, dim1=0, dim2=1)
+    assert t.shape == (SMALL_Y, LARGE_X)
+    assert np.sum(t[:,-1].asnumpy() == LARGE_X) == b.shape[1]
+    t = nd.flip(b, axis=0)
+    assert np.sum(t[-1,:].asnumpy() == 0) == b.shape[1]
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()