From aaa7f1c4a240546ecb4d396950f7e60496555059 Mon Sep 17 00:00:00 2001 From: Chaitanya Prakash Bapat Date: Thu, 7 Feb 2019 02:28:45 -0500 Subject: [PATCH 1/3] large op support --- src/operator/tensor/broadcast_reduce_op.h | 24 +++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/operator/tensor/broadcast_reduce_op.h b/src/operator/tensor/broadcast_reduce_op.h index 1edcb5a74a77..58f5c4145f02 100644 --- a/src/operator/tensor/broadcast_reduce_op.h +++ b/src/operator/tensor/broadcast_reduce_op.h @@ -1172,18 +1172,18 @@ void L2NormComputeEx(const nnvm::NodeAttrs& attrs, template struct pick { template - MSHADOW_XINLINE static void Map(int i, DType* out, const DType* a, - const IType *idx, int M, int stride, + MSHADOW_XINLINE static void Map(index_t i, DType* out, const DType* a, + const IType *idx, size_t M, int stride, mshadow::Shape bshape, mshadow::Shape sshape) { using namespace broadcast; - int j = static_cast(idx[i]); + index_t j = static_cast(idx[i]); if (clip) { if (j <= 0) j = 0; - else if (j >= M) j = M - 1; + else if (j >= static_cast(M)) j = static_cast(M) - 1; } else { - j = j % M; - j += (j < 0) ? M : 0; + j = j % static_cast(M); + j += (j < 0) ? static_cast(M) : 0; } j = ravel(unravel(i, sshape), bshape) + j*stride; out[i] = a[j]; @@ -1194,18 +1194,18 @@ struct pick { template struct pick_grad { template - MSHADOW_XINLINE static void Map(int i, DType* igrad, const DType* ograd, - const IType *idx, int M, int stride, + MSHADOW_XINLINE static void Map(index_t i, DType* igrad, const DType* ograd, + const IType *idx, size_t M, int stride, mshadow::Shape bshape, mshadow::Shape sshape) { using namespace broadcast; - int j = static_cast(idx[i]); + index_t j = static_cast(idx[i]); if (clip) { if (j <= 0) j = 0; - else if (j >= M) j = M - 1; + else if (j >= static_cast(M)) j = static_cast(M) - 1; } else { - j = j % M; - j += (j < 0) ? M : 0; + j = j % static_cast(M); + j += (j < 0) ? static_cast(M) : 0; } j = ravel(unravel(i, sshape), bshape) + j*stride; igrad[j] += ograd[i]; From 029850d37216878d37942998a7009b422ad38976 Mon Sep 17 00:00:00 2001 From: Chaitanya Prakash Bapat Date: Tue, 12 Feb 2019 17:44:39 -0500 Subject: [PATCH 2/3] replaced size_t with index_t for M, added test case --- src/operator/tensor/broadcast_reduce_op.h | 16 ++++++++-------- tests/nightly/test_large_array.py | 5 +++++ 2 files changed, 13 insertions(+), 8 deletions(-) diff --git a/src/operator/tensor/broadcast_reduce_op.h b/src/operator/tensor/broadcast_reduce_op.h index 58f5c4145f02..6aeeadfe820d 100644 --- a/src/operator/tensor/broadcast_reduce_op.h +++ b/src/operator/tensor/broadcast_reduce_op.h @@ -1173,17 +1173,17 @@ template struct pick { template MSHADOW_XINLINE static void Map(index_t i, DType* out, const DType* a, - const IType *idx, size_t M, int stride, + const IType *idx, index_t M, int stride, mshadow::Shape bshape, mshadow::Shape sshape) { using namespace broadcast; index_t j = static_cast(idx[i]); if (clip) { if (j <= 0) j = 0; - else if (j >= static_cast(M)) j = static_cast(M) - 1; + else if (j >= M) j = M - 1; } else { - j = j % static_cast(M); - j += (j < 0) ? static_cast(M) : 0; + j = j % M; + j += (j < 0) ? M : 0; } j = ravel(unravel(i, sshape), bshape) + j*stride; out[i] = a[j]; @@ -1195,17 +1195,17 @@ template struct pick_grad { template MSHADOW_XINLINE static void Map(index_t i, DType* igrad, const DType* ograd, - const IType *idx, size_t M, int stride, + const IType *idx, index_t M, int stride, mshadow::Shape bshape, mshadow::Shape sshape) { using namespace broadcast; index_t j = static_cast(idx[i]); if (clip) { if (j <= 0) j = 0; - else if (j >= static_cast(M)) j = static_cast(M) - 1; + else if (j >= M) j = M - 1; } else { - j = j % static_cast(M); - j += (j < 0) ? static_cast(M) : 0; + j = j % M; + j += (j < 0) ? M : 0; } j = ravel(unravel(i, sshape), bshape) + j*stride; igrad[j] += ograd[i]; diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index 696fdb1d4175..0e46a15e9212 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -145,6 +145,11 @@ def test_where(): res = nd.sparse.where(csr_cond, a, b) assert np.sum(res[0].asnumpy() == 1) == b.shape[1] +def test_pick(): + a = mx.nd.ones(shape=(LARGE_X, SMALL_Y)) + b = mx.nd.ones(shape=(LARGE_X,)) + res = mx.nd.pick(a,b) + assert res.shape == b.shape if __name__ == '__main__': import nose From f05959adfddce731d2ebb321f9bc8a1d3615e461 Mon Sep 17 00:00:00 2001 From: Chaitanya Prakash Bapat Date: Wed, 13 Feb 2019 22:55:23 -0500 Subject: [PATCH 3/3] changed shape --- tests/nightly/test_large_array.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index 0e46a15e9212..0249f44932c2 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -146,8 +146,8 @@ def test_where(): assert np.sum(res[0].asnumpy() == 1) == b.shape[1] def test_pick(): - a = mx.nd.ones(shape=(LARGE_X, SMALL_Y)) - b = mx.nd.ones(shape=(LARGE_X,)) + a = mx.nd.ones(shape=(256*35, 1024*1024)) + b = mx.nd.ones(shape=(256*35,)) res = mx.nd.pick(a,b) assert res.shape == b.shape