如何解包冒号与NumPy None列表以实现数组索引？

问题背景 #

我需要在n维NumPy数组的任意位置使用任意数量的:和None（即numpy.newaxis），想通过解包的方式把这些轴运算符放到数组索引的[]里，自动匹配对应轴。

Pylance提示：

下标中的解包运算符需要Python 3.11或更高版本

但升级到Python 3.11后，运行代码却报错：

Traceback (most recent call last):
  File "/home/.../quant.py", line 261, in <module>
    print(arr[*lhs_axes] + arr2[None,None,:])
          ~~~^^^^^^^^^^^
IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

当前代码：

import numpy as np
if __name__ == "__main__":
    lhs_ind, rhs_ind = 'ij', 'k'
        
    lhs_axes = [':' for i in lhs_ind]
    lhs_axes.append(None)

    arr1 = np.ones((2,2))
    arr2 = np.ones(2)
    print(arr1[*lhs_axes] + arr2[None,None,:])

解决方法 #

问题出在NumPy不支持直接解包列表作为索引，必须把存储索引运算符的容器改成元组。NumPy的多轴索引本质是接收元组参数，解包元组才能被正确识别为多个独立的轴索引。

修改后的代码：

import numpy as np
if __name__ == "__main__":
    lhs_ind, rhs_ind = 'ij', 'k'
        
    # 将列表转为元组
    lhs_axes = tuple([':' for _ in lhs_ind] + [None])

    arr1 = np.ones((2,2))
    arr2 = np.ones(2)
    # 解包元组进行索引
    print(arr1[*lhs_axes] + arr2[None,None,:])

更简洁的写法：

# 直接生成元组，避免列表转换
lhs_axes = tuple(':') * len(lhs_ind) + (None,)

原理说明 #

当你写arr[a, b, c]时，NumPy实际是把(a, b, c)这个元组作为索引参数。如果解包的是列表*lhs_axes，Python会把列表元素展开，但NumPy无法将这些展开后的元素识别为合法的多轴索引组合；而解包元组时，NumPy能正确解析每个元素对应的轴操作，不会触发索引错误。

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