在TensorFlow中根据指定张量重新排列张量的实现方法

不管你用NumPy、PyTorch还是TensorFlow，都能轻松实现c = a[b]这个按索引重排的需求，下面是各框架的具体示例：

NumPy 方案 #

NumPy原生支持整数数组索引，哪怕你的索引张量b是uint类型也完全没问题，只要满足max(b) < len(a)的条件，直接索引就行：

import numpy as np

# 示例张量a
a = np.array([10, 20, 30, 40, 50])
# uint32类型的索引张量b，max(b)=4 < len(a)=5
b = np.uint32([3, 1, 0, 4])
# 直接索引得到重排后的c
c = a[b]
print(c)  # 输出: [40 20 10 50]

NumPy会自动处理uint类型的索引值，不需要额外类型转换，只要索引在有效范围内就会返回正确的结果。

PyTorch 方案 #

PyTorch同样支持用uint类型张量作为索引，直接进行重排操作：

import torch

a = torch.tensor([10, 20, 30, 40, 50])
# uint8类型的索引张量
b = torch.tensor([3, 1, 0, 4], dtype=torch.uint8)
c = a[b]
print(c)  # 输出: tensor([40, 20, 10, 50])

如果遇到极少数类型兼容问题，也可以把b转换成长整型：c = a[b.long()]，结果完全一致。

TensorFlow 方案 #

在TensorFlow 2.x中，你既可以用直接索引的方式，也可以用tf.gatherAPI来实现，两种方式都支持uint类型索引：

import tensorflow as tf

a = tf.constant([10, 20, 30, 40, 50])
b = tf.constant([3, 1, 0, 4], dtype=tf.uint32)
# 方式1：直接索引
c1 = a[b]
# 方式2：使用tf.gather
c2 = tf.gather(a, b)
# 两种方式结果相同
print(c1.numpy())  # 输出: [40 20 10 50]
print(c2.numpy())  # 输出: [40 20 10 50]

tf.gather是TensorFlow专门用于按索引收集元素的工具，适合更复杂的索引场景，简单重排的话直接索引更直观。

总的来说，核心逻辑就是用索引张量b直接对a进行索引操作，主流张量库都原生支持uint类型的索引，只要保证b中的最大值不超过a的长度减一，就能得到你想要的重排结果。

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