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

Commit

Permalink
adding large tensor support for add_n and tests for more ops
Browse files Browse the repository at this point in the history
  • Loading branch information
Rohit Kumar Srivastava committed Oct 17, 2019
1 parent c2bbde7 commit 1917af8
Show file tree
Hide file tree
Showing 3 changed files with 112 additions and 4 deletions.
8 changes: 4 additions & 4 deletions src/operator/tensor/elemwise_sum.h
Original file line number Diff line number Diff line change
Expand Up @@ -39,15 +39,15 @@ namespace op {

struct Sum {
template<typename DType>
MSHADOW_XINLINE static DType sum(int i, const DType* a) {
MSHADOW_XINLINE static DType sum(index_t i, const DType* a) {
return a[i];
}
template<typename DType, typename... DTypes>
MSHADOW_XINLINE static DType sum(int i, const DType* a, const DTypes... b) {
MSHADOW_XINLINE static DType sum(index_t i, const DType* a, const DTypes... b) {
return a[i] + sum(i, b...);
}
template<typename DType, typename... DTypes>
MSHADOW_XINLINE static void Map(int i, DType* out, const OpReqType req, const DType* in0,
MSHADOW_XINLINE static void Map(index_t i, DType* out, const OpReqType req, const DType* in0,
const DTypes... ins) {
KERNEL_ASSIGN(out[i], req, sum(i, in0, ins...));
}
Expand All @@ -64,7 +64,7 @@ void ElementWiseSumCompute_(const nnvm::NodeAttrs& attrs,
size_t size = in_data.size();
Stream<xpu> *s = ctx.get_stream<xpu>();
DType* out_dptr = out_data[0].dptr<DType>();
int out_size = static_cast<int>((out_data[0].Size() + DataType<DType>::kLanes - 1)
index_t out_size = static_cast<index_t>((out_data[0].Size() + DataType<DType>::kLanes - 1)
/DataType<DType>::kLanes);
switch (size) {
case 2: {
Expand Down
54 changes: 54 additions & 0 deletions tests/nightly/test_large_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.

import os
import tempfile
import math
import numpy as np
import mxnet as mx
Expand Down Expand Up @@ -1199,6 +1201,58 @@ def test_full():
assert a[-1][-1] == 3


def test_load_save():
x = create_2d_tensor(SMALL_Y, LARGE_X)
tmp = tempfile.mkdtemp()
tmpfile = os.path.join(tmp, 'large_tensor')
nd.save(tmpfile, [x])
y = nd.load(tmpfile)
y = y[0]
assert x[0][0] == y[0][0]
assert x[-1][-1]== y[-1][-1]


def test_add_n():
x = [nd.ones(LARGE_X) for j in range(SMALL_Y)]
y = nd.add_n(*x)
assert y[0] == SMALL_Y
assert y[-1] == SMALL_Y


def test_modulo():
x = mx.nd.ones((SMALL_Y, LARGE_X))*6
y = mx.nd.ones(LARGE_X)*4
z = (x%y)
assert z[0][0] == 2
assert z[-1][-1] == 2
x = mx.nd.ones((SMALL_Y, LARGE_X))*5
z = nd.modulo(x,y)
assert z[0][0] == 1
assert z[-1][-1] == 1


def test_maximum():
x = mx.nd.ones((SMALL_Y, LARGE_X))*3
y = mx.nd.ones(LARGE_X)*4
z = nd.maximum(x, y)
assert z[0][0] == 4
assert z[-1][-1] == 4
z = nd.maximum(x, 5)
assert z[0][0] == 5
assert z[-1][-1] == 5


def test_minimum():
x = mx.nd.ones((SMALL_Y, LARGE_X))*3
y = mx.nd.ones(LARGE_X)*2
z = nd.minimum(x, y)
assert z[0][0] == 2
assert z[-1][-1] == 2
z = nd.minimum(x, 5)
assert z[0][0] == 3
assert z[-1][-1] == 3


if __name__ == '__main__':
import nose
nose.runmodule()
54 changes: 54 additions & 0 deletions tests/nightly/test_large_vector.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.

import os
import tempfile
import math
import numpy as np
import mxnet as mx
Expand Down Expand Up @@ -708,6 +710,58 @@ def test_full():
assert a[-1] == 3


def test_load_save():
x = create_vector(size=LARGE_X)
tmp = tempfile.mkdtemp()
tmpfile = os.path.join(tmp, 'large_vector')
nd.save(tmpfile, [x])
y = nd.load(tmpfile)
y = y[0]
assert x[0] == y[0]
assert x[-1] == y[-1]


def test_add_n():
x = [nd.ones(LARGE_X)]
y = nd.add_n(*x)
assert y[0] == 1
assert y[-1] == 1


def test_modulo():
x = mx.nd.ones(LARGE_X)*6
y = mx.nd.ones(LARGE_X)*4
z = (x%y)
assert z[0] == 2
assert z[-1] == 2
x = mx.nd.ones(LARGE_X)*5
z = nd.modulo(x,y)
assert z[0] == 1
assert z[-1] == 1


def test_maximum():
x = mx.nd.ones(LARGE_X)*3
y = mx.nd.ones(LARGE_X)*4
z = nd.maximum(x, y)
assert z[0] == 4
assert z[-1] == 4
z = nd.maximum(x, 5)
assert z[0] == 5
assert z[-1] == 5


def test_minimum():
x = mx.nd.ones(LARGE_X)*3
y = mx.nd.ones(LARGE_X)*2
z = nd.minimum(x, y)
assert z[0] == 2
assert z[-1] == 2
z = nd.minimum(x, 5)
assert z[0] == 3
assert z[-1] == 3


if __name__ == '__main__':
import nose
nose.runmodule()

0 comments on commit 1917af8

Please sign in to comment.