如何对前n个值进行排序，并将它们保存在metrix中，而其他值在Tensorflow中变为零？

import tensorflow as tf

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

sorted_indices = tf.argsort(input_tensor, direction='DESCENDING')

zero_matrix = tf.zeros_like(input_tensor)

n = 3  # 前n个值
sorted_values = tf.gather(input_tensor, sorted_indices[:n])
result_matrix = tf.tensor_scatter_nd_update(zero_matrix, tf.expand_dims(sorted_indices[:n], axis=1), sorted_values)