PyTorch张量索引官方文档位置查询及用法咨询

云原生阿岚

2026-7-28

PyTorch张量索引（t[idx]）详解及官方文档指引

核心操作逻辑

使用张量idx对t执行t[idx]属于高级索引操作，核心规则分以下几种典型场景：

当idx是一维整数张量：从t的第0维度（最外层维度）提取对应索引的元素，返回张量的形状与idx完全一致。例如：
```
import torch
t = torch.tensor([[1,2],[3,4],[5,6]])
idx = torch.tensor([0,2])
print(t[idx])  # 输出: tensor([[1, 2], [5, 6]])
```
当idx是多维整数张量：将idx中的每个元素作为t第0维度的索引，最终返回的张量形状和idx的形状一致。例如：
```
idx = torch.tensor([[0,1],[2,0]])
print(t[idx])  # 输出: tensor([[[1, 2], [3, 4]], [[5, 6], [1, 2]]])
```
多维度组合索引（如t[idx1, idx2]）：对t的多个维度分别应用索引，返回张量的形状为idx1与idx2广播后的形状。