如何在Tensorflow中计算矩阵对角线的和？

import tensorflow as tf

matrix = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=tf.float32)

diagonal = tf.linalg.diag_part(matrix)

diagonal_sum = tf.reduce_sum(diagonal)

import tensorflow as tf

matrix = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=tf.float32)
diagonal = tf.linalg.diag_part(matrix)
diagonal_sum = tf.reduce_sum(diagonal)

with tf.Session() as sess:
    result = sess.run(diagonal_sum)
    print("矩阵对角线的和为:", result)