有没有一种简单的方法来扩展现有的激活函数？我的自定义softmax函数返回:一个操作的渐变有`None`

import tensorflow as tf

def custom_softmax(x):
    # 自定义softmax函数
    exp_x = tf.exp(x)
    sum_exp_x = tf.reduce_sum(exp_x, axis=1, keepdims=True)
    softmax_x = exp_x / sum_exp_x
    return softmax_x

def custom_softmax_grad(op, grad):
    # 自定义softmax函数的梯度计算
    softmax_x = op.outputs[0]
    softmax_grad = grad * (softmax_x - tf.square(softmax_x))
    return softmax_grad

# 注册自定义softmax函数和梯度函数
@tf.RegisterGradient("CustomSoftmax")
def _custom_softmax_grad(op, grad):
    return custom_softmax_grad(op, grad)

# 使用自定义softmax函数
with tf.Session() as sess:
    x = tf.placeholder(tf.float32, shape=(None, 10))
    softmax = tf.nn.softmax(x)
    tf.register_gradient("CustomSoftmax")  # 注册自定义梯度函数
    softmax_grad = tf.gradients(softmax, x, grad_ys=tf.ones_like(softmax), name="CustomSoftmaxGrad")[0]

    # 测试自定义softmax函数和梯度函数
    input_data = [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0]]
    softmax_output, softmax_grad_output = sess.run([softmax, softmax_grad], feed_dict={x: input_data})
    print("Custom Softmax Output:", softmax_output)
    print("Custom Softmax Gradient:", softmax_grad_output)