Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit 2e1bd4e

Browse files
author
Rohit Kumar Srivastava
committed
[MXNET-1410]Adding Large Tensor Support for tensor transpose
1 parent 5fc4fc5 commit 2e1bd4e

File tree

2 files changed

+31
-4
lines changed

2 files changed

+31
-4
lines changed

src/operator/tensor/matrix_op-inl.h

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1950,10 +1950,10 @@ struct ReverseParam : public dmlc::Parameter<ReverseParam> {
19501950
#define REVERSE_MAX_DIM 10U
19511951

19521952
struct reverse {
1953-
MSHADOW_XINLINE static int ReverseIndex(index_t idx,
1954-
index_t nreversedim,
1955-
const index_t * stride_,
1956-
const index_t * trailing_) {
1953+
MSHADOW_XINLINE static index_t ReverseIndex(index_t idx,
1954+
index_t nreversedim,
1955+
const index_t * stride_,
1956+
const index_t * trailing_) {
19571957
index_t outputIndex = idx;
19581958
for (index_t i = 0; i < nreversedim; ++i) {
19591959
const index_t low = outputIndex % trailing_[i];

tests/nightly/test_large_array.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -279,6 +279,33 @@ def test_diag():
279279
assert_almost_equal(r.asnumpy(), np.diag(a_np, k=k))
280280

281281

282+
def create_2d_tensor(rows, columns):
283+
a = np.arange(0, rows).reshape(rows, 1)
284+
b = np.broadcast_to(a, shape=(a.shape[0], columns))
285+
return nd.array(b, dtype=np.int64)
286+
287+
288+
def test_transpose():
289+
b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y)
290+
t = b.T
291+
assert t.shape == (SMALL_Y, LARGE_X)
292+
assert np.sum(t[:, -1].asnumpy() == (LARGE_X - 1)) == b.shape[1]
293+
294+
295+
def test_swapaxes():
296+
b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y)
297+
t = nd.swapaxes(b, dim1=0, dim2=1)
298+
assert t.shape == (SMALL_Y, LARGE_X)
299+
assert np.sum(t[:, -1].asnumpy() == (LARGE_X - 1)) == b.shape[1]
300+
301+
302+
def test_flip():
303+
b = create_2d_tensor(rows=LARGE_X, columns=SMALL_Y)
304+
t = nd.flip(b, axis=0)
305+
assert t.shape == (LARGE_X, SMALL_Y)
306+
assert np.sum(t[-1, :].asnumpy() == 0) == b.shape[1]
307+
308+
282309
if __name__ == '__main__':
283310
import nose
284311
nose.runmodule()

0 commit comments

Comments
 (0)