Pytorch, efficient way extend a tensor by its first and last element
I have a tensor in pytorch. I want to extend it on a specific dimension from the beginning and the end by k positions with the first and last elements of that dimension respectively.
Say I have the tensor with data [[0, 0, 0], [1, 1, 1], [2, 2, 2]]. Operation extend(dim, k) would change it in this way:  

extend(0, 1):  [[0, 0, 0], [0, 0, 0], [0, 0, 0], [1, 1, 1], [2, 2, 2], [2, 2, 2], [2, 2, 2]]  

extend(1, 1): [0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [2, 2, 2, 2, 2]]  

What is an efficient way to do this (compliant with tensor.requires_grad=true)

You are looking for torch.nn.functional.pad, with mode='replicate'.
However, there are two things you need to pa attention to to get this to work:
1. pad does not work with 2D tensors. Thus, you need to add leading singleton dimensions before pad and squeezeing them afterwards.
2. The order of pad values pad expects is opposite to dim order.

import torch
from torch.nn inport functional
x = torch.tensor([[0, 0, 0],[1, 1, 1], [2, 2, 2]], dtype=torch.float)
# expand along dim=0 by k=2 
f.pad(x[None,None,...], (0,0, 2, 2), mode='replicate').squeeze()



Out[]:
tensor([[0., 0., 0.],
        [0., 0., 0.],
        [0., 0., 0.],
        [1., 1., 1.],
        [2., 2., 2.],
        [2., 2., 2.],
        [2., 2., 2.]])



# expand along dim=1 by k=2
f.pad(x[None,None,...], (2, 2, 0 , 0), mode='replicate').squeeze()



Out[]:
tensor([[0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 1., 1., 1., 1.],
        [2., 2., 2., 2., 2., 2., 2.]])


