Suppose I have a tensor 2D tensor x
of shape (n,m)
. How can I extend the first dimension of the tensor by appending zero rows in x
by specifying the indices of where the zero rows will be located in the resulting tensor? For a concrete example:
x = torch.tensor([[1,1,1],
[2,2,2],
[3,3,3],
[4,4,4]])
And I want to append 2 zero rows such that their row-index will be 1,3, respectively, in the resulting tensor? I.e. in the example the result would be
X = torch.tensor([1,1,1],
[0,0,0],
[2,2,2],
[0,0,0],
[3,3,3],
[4,4,4]])
I tried using F.pad
and reshape
.
You can use torch.tensor.index_add_
.
import torch
zero_index = [1, 3]
size = (6, 3)
x = torch.tensor([[1,1,1],
[2,2,2],
[3,3,3],
[4,4,4]])
t = torch.zeros(size, dtype=torch.int64)
index = torch.tensor([i for i in range(size[0]) if i not in zero_index])
# index -> tensor([0, 2, 4, 5])
t.index_add_(0, index, x)
print(t)
Output:
tensor([[1, 1, 1],
[0, 0, 0],
[2, 2, 2],
[0, 0, 0],
[3, 3, 3],
[4, 4, 4]])