如何用NumPy实现任意批量维度的向量列表批量索引？

给定维度为(N1,...,Nk,L,H)的数组data，以及维度为(N1,...,Nk)的批量索引数组inds，我们需要直接用NumPy操作替代循环，得到维度为(N1,...,Nk,H)的输出——核心是为每个批量位置(i1,...,ik)选取data[i1,...,ik, inds[i1,...,ik], :]对应的H维向量。

方法一：用np.ix_生成网格索引 #

np.ix_可以生成多维网格索引，完美匹配批量维度的索引需求，再结合inds即可直接完成索引：

import numpy as np

# 测试案例1：三维批量维度
N1 = (3, 3, 3)
L, H = 5, 2
data = np.arange(np.prod(N1) * L * H).reshape(*N1, L, H)
inds = np.arange(np.prod(N1), dtype=int).reshape(N1) % L

# 核心操作
batch_dims = data.shape[:-2]
grid = np.ix_(*[np.arange(d) for d in batch_dims]) if batch_dims else ()
out = data[grid + (inds,)]

assert out.shape == (*N1, H)

# 测试案例2：空批量维度
N2 = ()
data = np.arange(np.prod(N2) * L * H).reshape(*N2, L, H)
inds = np.arange(np.prod(N2), dtype=int).reshape(N2) % L

batch_dims = data.shape[:-2]
grid = np.ix_(*[np.arange(d) for d in batch_dims]) if batch_dims else ()
out = data[grid + (inds,)]

assert out.shape == (*N2, H)

方法二：用np.take_along_axis沿轴选取 #

take_along_axis专门用于沿指定轴根据索引选取元素，只需将inds扩展维度匹配data的结构即可：

import numpy as np

# 测试案例1：三维批量维度
N1 = (3, 3, 3)
L, H = 5, 2
data = np.arange(np.prod(N1) * L * H).reshape(*N1, L, H)
inds = np.arange(np.prod(N1), dtype=int).reshape(N1) % L

# 核心操作
inds_expanded = np.expand_dims(inds, axis=(-1, -2))  # 扩展为(*batch_dims, 1, 1)
out = np.take_along_axis(data, inds_expanded, axis=-2).squeeze(axis=-2)

assert out.shape == (*N1, H)

# 测试案例2：空批量维度
N2 = ()
data = np.arange(np.prod(N2) * L * H).reshape(*N2, L, H)
inds = np.arange(np.prod(N2), dtype=int).reshape(N2) % L

inds_expanded = np.expand_dims(inds, axis=(-1, -2))
out = np.take_along_axis(data, inds_expanded, axis=-2).squeeze(axis=-2)

assert out.shape == (*N2, H)

两种方法都能完全替代原循环，且不需要展平批量维度，直接利用NumPy的高级索引机制实现高效操作。

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