Skip to content

Commit

Permalink
[MXNET-1400] adding tests cases to verify large tensor support for de…
Browse files Browse the repository at this point in the history
…pth_to_space and space_to_depth (apache#14797)
  • Loading branch information
access2rohit authored and haohuw committed Jun 23, 2019
1 parent 2764f07 commit 7e14b79
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 1 deletion.
1 change: 1 addition & 0 deletions CONTRIBUTORS.md
Original file line number Diff line number Diff line change
Expand Up @@ -239,6 +239,7 @@ List of Contributors
* [Zak Jost](https://github.com/zjost)
* [Shoubhik Bhattacharya](https://github.com/shoubhik)
* [Zach Kimberg](https://github.com/zachgk)
* [Rohit Srivastava](https://github.com/access2rohit)

Label Bot
---------
Expand Down
32 changes: 31 additions & 1 deletion tests/nightly/test_large_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

import mxnet as mx
import numpy as np
from mxnet.test_utils import rand_ndarray, assert_almost_equal
from mxnet import gluon, nd
from tests.python.unittest.common import with_seed

Expand Down Expand Up @@ -185,7 +186,36 @@ def test_pick():
b = mx.nd.ones(shape=(256*35,))
res = mx.nd.pick(a,b)
assert res.shape == b.shape


def test_depthtospace():
def numpy_depth_to_space(x, blocksize):
b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
tmp = np.reshape(x, [b, blocksize, blocksize, c // (blocksize**2), h, w])
tmp = np.transpose(tmp, [0, 3, 4, 1, 5, 2])
y = np.reshape(tmp, [b, c // (blocksize**2), h * blocksize, w * blocksize])
return y

shape_inp = (LARGE_X, 8, 4, 2)
data = rand_ndarray(shape_inp, 'default')
data_np = data.asnumpy()
expected = numpy_depth_to_space(data_np, 2)
output = mx.nd.depth_to_space(data, 2)
assert_almost_equal(output.asnumpy(), expected, atol=1e-3, rtol=1e-3)

def test_spacetodepth():
def numpy_space_to_depth(x, blocksize):
b, c, h, w = x.shape[0], x.shape[1], x.shape[2], x.shape[3]
tmp = np.reshape(x, [b, c, h // blocksize, blocksize, w // blocksize, blocksize])
tmp = np.transpose(tmp, [0, 3, 5, 1, 2, 4])
y = np.reshape(tmp, [b, c * (blocksize**2), h // blocksize, w // blocksize])
return y

shape_inp = (LARGE_X, 2, 8, 4)
data = rand_ndarray(shape_inp, 'default')
data_np = data.asnumpy()
expected = numpy_space_to_depth(data_np, 2)
output = mx.nd.space_to_depth(data, 2)
assert_almost_equal(output.asnumpy(), expected, atol=1e-3, rtol=1e-3)

if __name__ == '__main__':
import nose
Expand Down

0 comments on commit 7e14b79

Please sign in to comment.