有一个Softmax激活函数和tf.nn.softmax（），当激活需要一个支持图形的层时，你可以选择它，它可以实现到一个模型中，但需要硬件支持。
示例：实现简单的自定义SoftMax层，您需要做它的右轴水平。

import tensorflow as tf

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
None
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
config = tf.config.experimental.set_memory_growth(physical_devices[0], True)
print(physical_devices)
print(config)

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Class / Funtions
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
class B_Softmax(tf.keras.layers.Layer):
    def __init__(self, units):
        super(B_Softmax, self).__init__()
        self.units = units
        self._out_shape = None

    def build(self, input_shape):
        self._out_shape = input_shape

    def call(self, inputs):
        temp = tf.transpose(inputs)
        temp = tf.keras.layers.Dense(self.units, activation=tf.nn.softmax)(temp)
        temp = tf.transpose(inputs)
        return temp

"""""""""""""""""""""""""""""""""""""""""""""""""""""""""
: Tasks
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""       
temp = tf.constant([[ 0.00346701, -0.00676209, -0.00109781, -0.0005832 ,  0.00047849, 0.00311204,  0.00843922, -0.00400238,  0.00127922, -0.0026469 ,
-0.00232184, -0.00686269,  0.00021552, -0.0039388 ,  0.00753652,
-0.00405236, -0.0008759 ,  0.00275771,  0.00144688, -0.00361056,
-0.0036177 ,  0.00778807, -0.00116923,  0.00012773,  0.00276652,
0.00438983, -0.00769166, -0.00432891, -0.00211244, -0.00594028,
0.01009954,  0.00581804, -0.0062736 , -0.00921499,  0.00710281,
0.00022364,  0.00051054, -0.00204145,  0.00928543, -0.00129213,
-0.00209933, -0.00212295, -0.00452125, -0.00601313, -0.00239222,
0.00663724,  0.00228883,  0.00359715,  0.00090024,  0.01166699,
-0.00281386, -0.00791688,  0.00055902,  0.00070648,  0.00052972,
0.00249906,  0.00491098,  0.00528313, -0.01159694, -0.00370812,
-0.00950641,  0.00408999,  0.00800613,  0.0014898 ]], dtype=tf.float32)
#  shape=(64, 10), dtype=float32

layer = B_Softmax(10)
print( layer( temp ) )
#  shape=(64, 1), dtype=float32

输出：SoftMax输出自定义反馈的简单输入。

tf.Tensor(
[[0.10015144 0.1000239  0.10018992 0.0999647  0.10004678 0.09998975
  0.09980123 0.09972709 0.10010113 0.10000402]
...
 [0.10034979 0.10005493 0.10043884 0.09991822 0.1001078  0.09997606
  0.09954134 0.09937066 0.10023339 0.10000902]
 [0.10006507 0.10001029 0.10008159 0.09998485 0.10002013 0.09999561
  0.09991457 0.09988266 0.10004347 0.10000175]], shape=(64, 10), dtype=float32)

tf.Tensor(
[[ 0.00346701]
 [-0.00676209]
...
 [ 0.00408999]
 [ 0.00800613]
 [ 0.0014898 ]], shape=(64, 1), dtype=float32)