神经网络中的Softmax函数(Python)

Question

我正在学习神经网络，并在python中实现它。我首先定义了一个 softmax 函数，我遵循这个问题给出的解决方案 Softmax 函数 - python。以下是我的代码：

def softmax(A):
    """
    Computes a softmax function. 
    Input: A (N, k) ndarray.
    Returns: (N, k) ndarray.
    """
    s = 0
    e = np.exp(A)
    s = e / np.sum(e, axis =0)
    return s

我得到了一个测试代码，看看sofmax函数是否正确。test_array是测试数据，test_output的正确输出。以下是测试代码：

# Test if your function works correctly.
test_array = np.array([[0.101,0.202,0.303],
                       [0.404,0.505,0.606]]) 
test_output = [[ 0.30028906,  0.33220277,  0.36750817],
               [ 0.30028906,  0.33220277,  0.36750817]]
print(np.allclose(softmax(test_array),test_output))

但是根据我定义的softmax函数。通过softmax（test_array）测试数据返回

print (softmax(test_array))

[[ 0.42482427  0.42482427  0.42482427]
 [ 0.57517573  0.57517573  0.57517573]]

有人能告诉我我定义的函数softmax有什么问题吗？