pythonpytorchtensorzero-padding

Appending zero rows to a 2D Tensor in PyTorch


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.


Solution

  • 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]])