pytorch 的 index_add 操作

 

基础操作¶

 

In [16]:

import torch
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"

# 目标张量 (shape: [3, 4])
input = torch.zeros(3, 4,dtype=torch.long)

# 源张量 (shape: [2, 4])
source = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])

# 索引张量 (shape: [2])
index = torch.tensor([0, 2])  # 将source的第0行加到input的第0行，第1行加到input的第2行

# 沿dim=0操作
result = input.index_add(0, index, source)

print(result)

 

 

tensor([[1, 2, 3, 4],
        [0, 0, 0, 0],
        [5, 6, 7, 8]])

 

 

不同维度的操作¶

 

In [15]:

# 目标张量 (shape: [4, 3])
input = torch.zeros(4, 3,dtype=torch.long)

# 源张量 (shape: [4, 2])
source = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])

# 索引张量 (shape: [2])
index = torch.tensor([1, 2])  # 将source的列0加到input的列1，列1加到input的列2

# 沿dim=1操作
result = input.index_add(1, index, source)

print(result)

 

 

tensor([[0, 1, 2],
        [0, 3, 4],
        [0, 5, 6],
        [0, 7, 8]])