numpy softmax广播算法练习

def softmax(x):
    x_exp = np.exp(x)
    print('x_exp: ', x_exp.shape)
    ## 输出为(2, 5) 和本来的shape一致
    x_sum = np.sum(x_exp, axis=1, keepdims=True)
    ## keepdims非常重要不然的话x_sum变成rank1 array之后少了一个[]无法进行广播计算
    print('x_sum', x_sum.shape)
    ## (输出为2, 1)
    softmax = x_exp/x_sum
    return softmax
x = np.array([
    [9, 2, 5, 0, 0],
    [7, 5, 0, 0 ,0]])
print("softmax(x) = " + str(softmax(x)))