如何用NumPy实现批量向量与批量索引矩阵的索引？支持高维批量吗？

一、针对给定场景的最简实现 #

对于你给出的批量向量vecs(形状(3,4))和批量索引矩阵mats(形状(3,2,2))，直接用NumPy的高级广播索引就能完成，代码非常简洁：

import numpy as np

vecs = np.asarray([[1, 2, 3, 4],
                   [5, 6, 7, 8],
                   [9,10,11,12]])
mats = np.asarray([ [[0,1], [1,0]],
                    [[0,2], [1,1]],
                    [[3,1], [0,0]] ])

# 生成批次索引并广播匹配mats的维度
batch_indices = np.arange(vecs.shape[0])[:, np.newaxis, np.newaxis]
results = vecs[batch_indices, mats]

print(results)
# 输出：
# [[[ 1  2]
#   [ 2  1]]
# 
#  [[ 5  7]
#   [ 6  6]]
# 
#  [[12 10]
#   [ 9  9]]]

原理：batch_indices把初始的(3,)形状扩展为(3,1,1)，和mats的(3,2,2)通过广播对齐，确保每个批次的向量只会被对应批次的索引矩阵索引，最终得到预期的(3,2,2)结果。

二、高维度批量的通用向量化方法 #

如果你的批量维度更多（比如vecs是(B1,B2,B3,D)，mats是(B1,B2,B3,I1,I2,I3)），可以用NumPy的np.take_along_axis函数，它专门用于沿指定轴提取对应索引的元素，无需手动构造复杂的批量索引：

# 通用实现：适配任意多的批量维度
# 步骤1：将vecs扩展为和mats维度一致，特征轴保持不变
vecs_expanded = np.expand_dims(vecs, axis=tuple(range(1, len(mats.shape) - len(vecs.shape) + 2)))
# 步骤2：将mats扩展出特征轴对应的维度，确保和vecs_expanded对齐
mats_expanded = np.expand_dims(mats, axis=1)
# 步骤3：沿特征轴（axis=1）提取元素
results = np.take_along_axis(vecs_expanded, mats_expanded, axis=1)
# 去掉多余的维度（可选，和预期结果维度匹配）
results = np.squeeze(results, axis=1)

针对你给出的场景，也可以用简化版的take_along_axis用法：

results = np.take_along_axis(vecs[:, :, np.newaxis, np.newaxis], mats[:, np.newaxis, :, :], axis=1)
results = np.squeeze(results, axis=1)

原理：take_along_axis会自动处理广播对齐，只要vecs和mats的前N个批量维度完全匹配，不管维度数量多少，都能正确完成每个批量内部的索引操作，避免手动构造索引的繁琐。

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