如何在tensorflow中定义此op的梯度函数？

import tensorflow as tf

def my_op(x):
    # 定义自定义操作
    return tf.square(x)

def my_op_grad(op, grad):
    # 定义梯度函数
    x = op.inputs[0]
    return [2 * x * grad]

# 注册自定义操作的梯度函数
@tf.RegisterGradient("MyOp")
def _my_op_grad(op, grad):
    return my_op_grad(op, grad)

# 使用自定义操作
with tf.GradientTape() as tape:
    x = tf.constant(2.0)
    y = my_op(x)

# 计算梯度
grad = tape.gradient(y, x)
print(grad)  # 输出：tf.Tensor(8.0, shape=(), dtype=float32)