pytorch 笔记：张量索引的维度扩展规则

1 基本原理

在PyTorch中，张量索引的维度扩展规则遵循以下原则：

索引操作的核心规则：​
当使用索引数组访问张量时：

1. 索引数组的每个元素对应选取原张量的一个子张量
2. 结果形状 = 索引数组形状 + 原张量剩余维度形状

这么说可能不清楚，我们用例子说明

2 举例

假如我们有这样的两个张量

mlp_weights.shape = (num_poi, mer_size, mer_dim)
top2_indices.shape = (batch, cur_seq_len, 2)

selected_weights = mlp_weights[top2_indices]

步骤	操作	维度变化说明
1	索引数组作用在维度0	原张量第0维度(num_poi)被索引数组替换
2	索引数组形状保持	继承索引数组的(batch, cur_seq_len, 2)形状
3	附加后续维度	保留原张量的(mer_size, mer_dim)维度

——>selected_weights的维度是 (batch, cur_seq_len, 2,mer_size, mer_dim)

2.1 索引过程

1. 每个索引值对应一个(mer_size,mer_dim)的矩阵：

   • top2_indices[0,0,0] = 1 → 选取mlp_weights[1,:,:]
   • top2_indices[0,0,1] = 2 → 选取mlp_weights[2,:,:]