From a37cd7abd547dbe27a7a370208065244ee39a35a Mon Sep 17 00:00:00 2001 From: Rohit Kumar Srivastava Date: Mon, 3 Jun 2019 21:32:24 -0700 Subject: [PATCH] [MXNET-1405] tests for large tensor support for Softmax operator (#15042) --- tests/nightly/test_large_array.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/nightly/test_large_array.py b/tests/nightly/test_large_array.py index 8c80e9e22488..cbba608d5d2f 100644 --- a/tests/nightly/test_large_array.py +++ b/tests/nightly/test_large_array.py @@ -319,6 +319,13 @@ def test_flip(): assert np.sum(t[-1, :].asnumpy() == 0) == b.shape[1] +def test_softmax(): + input_data = mx.nd.ones((SMALL_Y, LARGE_X)) + true_output = np.full((SMALL_Y, LARGE_X), (1 / SMALL_Y)) + output = nd.softmax(input_data, axis=0) + assert_almost_equal(output.asnumpy(), true_output, rtol=1e-5, atol=1e-5) + + if __name__ == '__main__': import nose nose.runmodule()