python实现softmax函数

import numpy as np

# 定义softmax函数，softmax:输出0-1之间的数，其和为1，可以用来表征判定是某个类别的概率
def softmax(x):
    
    # x为一维数据时
    if x.ndim == 1:
        return np.exp(x-np.max(x))/np.sum(np.exp(x-np.max(x)))
    
    # x为二维数据时
    elif x.ndim == 2:
        val_num = np.zeros_like(x)
        for i in range(len(x)):
            part_x = x[i]
            val_num[i] = np.exp(part_x-np.max(part_x))/np.sum(np.exp(part_x-np.max(part_x)))
        return val_num

softmax(np.array([1,3,5]))

array([0.01587624, 0.11731043, 0.86681333])

softmax(np.array([[1.0,3,5],[2,4,6],[3,6,9]]))

array([[0.01587624, 0.11731043, 0.86681333],
       [0.01587624, 0.11731043, 0.86681333],
       [0.00235563, 0.04731416, 0.95033021]])