如何在Tensorflow中编写一个简单的循环(对于矩阵的元素)？

import tensorflow as tf

# 定义循环条件
def loop_condition(i, matrix):
    return tf.less(i, tf.size(matrix))

# 定义循环体操作
def loop_body(i, matrix):
    # 获取矩阵元素
    element = tf.gather(tf.reshape(matrix, [-1]), i)
    
    # 对元素进行操作
    element_square = tf.square(element)
    
    # 更新矩阵元素
    matrix = tf.tensor_scatter_nd_update(matrix, [[i]], [element_square])
    
    # 循环变量自增
    i = tf.add(i, 1)
    
    return i, matrix

# 创建输入矩阵
matrix = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])

# 初始化循环变量
i = tf.constant(0)

# 执行循环
i, result_matrix = tf.while_loop(loop_condition, loop_body, loop_vars=[i, matrix])

# 打印结果
with tf.Session() as sess:
    result = sess.run(result_matrix)
    print(result)