Scatter tensor in pytorch along the rows

Viewed 599

I want to scatter tensors in granularities of rows.

For example consider,

Input = torch.tensor([[2, 3], [3, 4], [4, 5]])

I want to scatter

S = torch.tensor([[1,2],[1,2]])

to indices

I = torch.tensor([0,2])

I expect the output to be torch.tensor([[1, 2], [3, 4], [1, 2]]).

Here S[0] is scattered to Input[I[0]], similarly S[1] is scattered to Input[I[1]]

How can I achieve this? Instead of looping over the row in S, I am looking for a more efficient way.

2 Answers

Do input[I] = S

Example:

input = torch.tensor([[2, 3], [3, 4], [4, 5]])
S = torch.tensor([[1,2],[1,2]])
I = torch.tensor([0,2])
input[I] = S
input
tensor([[1, 2],
        [3, 4],
        [1, 2]])

Answer might be a little late, but anyway, you could do:

import torch

inp = torch.tensor([[2,3], [3,4], [4,5]])
src = torch.tensor([[1,2], [1,2]])
idxs = torch.tensor([[0,0],[2,2]])

y = torch.scatter(inp, 0, idxs, src)
Related