为何Numpy数组索引对(1,)形状ndarray与PyTorch tensor表现不同？

核心差异 #

• Numpy用(1,) ndarray索引：触发高级索引机制，会保留索引数组的维度，返回形状为(1,)的一维数组（如array([2])）。
• PyTorch用(1,) tensor索引Numpy数组：Numpy会自动将该tensor转换为Python标量（等价于用单个整数索引），触发基本索引，返回原数组的单个元素（如2）。

原因解析 #

1. Numpy高级索引特性：当索引是任意非tuple的序列类（比如ndarray）时，Numpy会启用高级索引。这种索引的核心是“索引的维度会被保留到结果中”，所以形状为(1,)的索引数组会让结果也保持(1,)的形状。
2. PyTorch tensor与Numpy的交互规则：PyTorch的tensor作为Numpy数组的索引时，Numpy会尝试将其转换为Python原生类型。对于形状为(1,)的tensor，Numpy会提取其唯一元素作为标量使用，此时就相当于用整数1去索引Numpy数组，自然返回单个标量值。

验证代码 #

import numpy as np
import torch as th

x = np.arange(10)

# 验证tensor转标量
print(th.tensor([1]).item())  # 输出：1（Python整数）
# 用整数索引对比
print(x[1])  # 输出：2，和z的结果完全一致

内容的提问来源于stack exchange，提问作者Nicolás_Tsu