如何在Tensorflow张量中选择每行和每列的前Q个元素？

import tensorflow as tf

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

Q = 3  # 前3个元素

top_k_values, top_k_indices = tf.math.top_k(x, Q, sorted=True, axis=1)

print(top_k_values)
print(top_k_indices)

tf.Tensor(
[[ 5  4  3]
 [10  9  8]], shape=(2, 3), dtype=int32)

tf.Tensor(
[[3 1 4]
 [4 1 3]], shape=(2, 3), dtype=int32)

import tensorflow as tf

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

Q = 3  # 前3个元素

top_k_values, top_k_indices = tf.math.top_k(x, Q, sorted=True, axis=0)

print(top_k_values)
print(top_k_indices)

tf.Tensor(
[[ 6  9  7  8 10]
 [ 4  4  3  5  6]
 [ 2  2  2  2  3]], shape=(3, 5), dtype=int32)

tf.Tensor(
[[1 1 1 1 1]
 [0 0 0 0 0]
 [0 0 0 0 0]], shape=(3, 5), dtype=int32)