如何在TensorFlow中基于两列对张量进行排序？

当然可以在TensorFlow里实现基于多列的张量排序！我给你两种实用的方法，结合你的例子来演示，保证能得到你想要的结果：

方法一：使用tf.lexsort（最直接的字典序排序） #

tf.lexsort是TensorFlow专门用来处理多列字典序排序的工具，它会按照你传入的键列表从右到左依次排序，正好适配「先按第一列、再按第二列」的需求：

import tensorflow as tf

# 定义你给出的原始张量
tensor = tf.constant([[1,2,3], [2,3,5], [1,4,6], [2,2,1], [0,4,2]])

# 提取排序键：注意顺序是「次要键在前，主键在后」
# 我们要先按第一列（索引0）排序，再按第二列（索引1）排序，所以次要键是第二列，主键是第一列
sort_keys = [tensor[:, 1], tensor[:, 0]]
# 获取排序后的索引
sorted_indices = tf.lexsort(sort_keys)
# 根据索引重新排列原张量
sorted_tensor = tf.gather(tensor, sorted_indices)

# 打印结果
print(sorted_tensor.numpy())
# 输出：
# [[0 4 2]
#  [1 2 3]
#  [1 4 6]
#  [2 2 1]
#  [2 3 5]]

简单解释下：tf.lexsort会从右往左处理你传入的键列表，所以把「第二列」（次要排序条件）放在前面，「第一列」（主要排序条件）放在后面，就会先按第一列排序，第一列值相同的元素再按第二列排序，完美匹配你的要求。

方法二：两次稳定排序（分步更直观） #

如果想更清晰地看到排序的分步过程，可以用两次稳定排序：先按次要键排序，再对主键做稳定排序（稳定排序会保留相同主键元素的相对顺序，也就是之前次要键的排序结果）：

import tensorflow as tf

tensor = tf.constant([[1,2,3], [2,3,5], [1,4,6], [2,2,1], [0,4,2]])

# 第一步：先按第二列（次要键）排序
sorted_by_col2 = tf.gather(tensor, tf.argsort(tensor[:, 1], stable=True))
# 第二步：按第一列（主键）做稳定排序，保证第一列相同的元素保留第二列的排序结果
sorted_final = tf.gather(sorted_by_col2, tf.argsort(sorted_by_col2[:, 0], stable=True))

print(sorted_final.numpy())
# 输出和方法一完全一致

这种方法的好处是逻辑更直观，每一步都能看到中间结果，适合理解多列排序的底层逻辑。

内容的提问来源于stack exchange，提问作者SameeraR