diff --git a/tests/python/mkl/test_mkldnn.py b/tests/python/mkl/test_mkldnn.py index e5490096f60f..c6c0a0832f1f 100644 --- a/tests/python/mkl/test_mkldnn.py +++ b/tests/python/mkl/test_mkldnn.py @@ -383,6 +383,21 @@ def check_fullyconnected_training(stype): for stype in stypes: check_fullyconnected_training(stype) +def test_softmax_with_large_inputs(): + def softmax_forward(input_data, true_output): + data = mx.sym.Variable('data') + out1 = data.softmax(axis=1) + exec1 = out1.bind(mx.cpu(), args={'data': input_data}) + exec1.forward()[0].wait_to_read() + ndarr = exec1.outputs[0][0][0][0] + nparr = ndarr.asnumpy() + assert_almost_equal(nparr, true_output, rtol=1e-5, atol=1e-5) + + softmax_forward(mx.nd.array([[[[-1e30,-1e30]]]]), np.array([1.0,1.0])) + softmax_forward(mx.nd.array([[[[1e30,1e30]]]]), np.array([1.0,1.0])) + softmax_forward(mx.nd.array([[[[-3.4e38,-3.4e38]]]]), np.array([1.0,1.0])) + softmax_forward(mx.nd.array([[[[3.4e38,3.4e38]]]]), np.array([1.0,1.0])) + @with_seed() def test_non_mkldnn_fcomputeex(): # test special case where MKLDNN formatted NDArray feeds into non-mkldnn fcomputeex operator diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index 1bf9ca0237ab..19b90b09ceb5 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -4449,6 +4449,20 @@ def test_log_softmax(): check_symbolic_forward(sym, [data], [np.log(np_softmax(data, axis=axis)+1e-20)]) check_numeric_gradient(sym, [data], rtol=0.05, atol=1e-3) +def test_softmax_with_large_inputs(): + def softmax_forward(input_data, true_output): + data = mx.sym.Variable('data') + out1 = data.softmax(axis=1) + exec1 = out1.bind(default_context(), args={'data': input_data}) + exec1.forward()[0].wait_to_read() + ndarr = exec1.outputs[0][0][0][0] + nparr = ndarr.asnumpy() + assert_almost_equal(nparr, true_output, rtol=1e-5, atol=1e-5) + + softmax_forward(mx.nd.array([[[[-1e30,-1e30]]]]), np.array([1.0,1.0])) + softmax_forward(mx.nd.array([[[[1e30,1e30]]]]), np.array([1.0,1.0])) + softmax_forward(mx.nd.array([[[[-3.4e38,-3.4e38]]]]), np.array([1.0,1.0])) + softmax_forward(mx.nd.array([[[[3.4e38,3.4e38]]]]), np.array([1.0,1.0])) @with_seed() def test_pick():