[MXNET-1405] tests for large tensor support for Softmax operator (#15042

)
apache · Jun 4, 2019 · a37cd7a · a37cd7a
1 parent 3ebd873
commit a37cd7a
Showing 1 changed file with 7 additions and 0 deletions.
diff --git 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()