TensorFlow如何实现PyTorch式的n维张量高级索引？

方法一：使用tf.gather_nd（推荐） #

这是最直接对应PyTorch多维度索引的方式，核心是把各维度的索引组合成统一的索引张量：

1. 将一维索引张量a、b、c、d沿第二个维度堆叠，生成形状为[9, 4]的索引矩阵，每一行对应前4个维度的一组坐标：

indices = tf.stack([a, b, c, d], axis=1)

2. 调用tf.gather_nd从preds中取值，它会根据索引矩阵里的每组坐标，取出对应位置的整个第五维（长度12）数据：

result = tf.gather_nd(preds, indices)

最终result的形状就是[9, 12]，和PyTorch的结果完全一致。

方法二：嵌套使用tf.gather #

虽然tf.gather本身是单维度索引，但通过嵌套调用并配合batch_dims参数，也能实现相同效果：

# 第0维索引，得到形状[9, 3, 64, 64, 12]
temp1 = tf.gather(preds, a, axis=0)
# 第1维索引，指定batch_dims=1让每个batch元素对应自身索引，得到[9, 64, 64, 12]
temp2 = tf.gather(temp1, b, axis=1, batch_dims=1)
# 继续第1维索引，得到[9, 64, 12]
temp3 = tf.gather(temp2, c, axis=1, batch_dims=1)
# 最后第1维索引，得到[9, 12]
result = tf.gather(temp3, d, axis=1, batch_dims=1)

这种方法步骤较多，不如tf.gather_nd简洁直观，仅作为扩展思路。

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