如何在TensorFlow 2.11中基于第四维度首元素排序4D张量？

问题说明 #

给定形状为(1, 6, 4, 2)的4D张量，要求保持第四维度内元素顺序不变，仅基于每个第三维度元素的第四维度首个值，对第三维度的数组进行升序排序。

解决方案代码 #

import tensorflow as tf

# 输入张量
input_tensor = tf.constant([[[[51, 92],
                              [14, 71],
                              [60, 20],
                              [82, 86]],
                             [[74, 74],
                              [87, 99],
                              [23,  2],
                              [21, 52]],
                             [[ 1, 87],
                              [29, 37],
                              [ 1, 63],
                              [59, 20]],
                             [[32, 75],
                              [57, 21],
                              [88, 48],
                              [90, 58]],
                             [[41, 91],
                              [59, 79],
                              [14, 61],
                              [61, 46]],
                             [[61, 50],
                              [54, 63],
                              [ 2, 50],
                              [ 6, 20]]]], dtype=tf.int64)

# 1. 提取排序依据：第三维每个元素的第四维首个值
sort_keys = input_tensor[..., 0]

# 2. 获取第三维的排序索引（升序）
sort_indices = tf.argsort(sort_keys, axis=2)

# 3. 按索引重新排列第三维元素，保持其他维度不变
sorted_tensor = tf.gather(input_tensor, sort_indices, axis=2, batch_dims=2)

# 验证结果
print(sorted_tensor.numpy())

代码解释 #

1. 提取排序键：通过input_tensor[..., 0]取出每个第三维度元素的第四维度第一个值，得到形状为(1, 6, 4)的排序依据张量。
2. 生成排序索引：使用tf.argsort在第三维（axis=2）上计算升序排序的索引，索引张量形状与排序键一致，每个值对应原第三维中元素的位置。
3. 重排张量：使用tf.gather指定axis=2（目标排序维度），并设置batch_dims=2，表示前两维（batch和第二维）作为固定的批次维度，仅对第三维按索引重新排列，确保第四维度内的元素顺序完全保留。

关于之前尝试的说明 #

如果之前使用tf.gather_nd未成功，通常是因为需要构造包含前两维索引的复杂多维索引张量，而tf.gather配合batch_dims参数更适合这种沿单一维度排序且保持其他维度结构的场景，代码更简洁且不易出错。

内容的提问来源于stack exchange，提问作者Soli Technology LLC