Implementation of function like `numpy.roll`

Ja-Keoung_Koo: Ultimately, I want to implement image gradient with forward difference and Neumann boundary condition. For example, the numpy version of it is as follows: <SCODE>def grad(u):
    # u: 2-d images
    ux = np.roll(u, -1, axis=1) - u
    uy = np.roll(u, -1, axis=0) - u
    ux[:,-1] = 0
    uy[-1,:] = 0
<ECODE>

Ja-Keoung_Koo: Thanks for your reply, although I checked before. As you said, it seems hard to implement.

apaszke: If should be quite simple to implement yourself. Just slice the tensor into two pieces, swap them, and cat along the same dimension that you used to split.

Ryan_Burgert: 

azban: haven’t tested this extensively, but this seems to cover if you just want a single split. logic if shift is negative could probably be cleaned up a little <SCODE>def roll(tensor, shift, axis):
    if shift == 0:
        return tensor

    if axis < 0:
        axis += tensor.dim()

    dim_size = tensor.size(axis)
    after_start = dim_size - shift
    if shift < 0:
        after_start = -shift
        shift = dim_size - abs(shift)

    before = tensor.narrow(axis, 0, dim_size - shift)
    after = tensor.narrow(axis, after_start, shift)
    return torch.cat([after, before], axis)
<ECODE>

jaromiru: Simple solution to roll around first axis: <SCODE>def roll(x, n):  
    return torch.cat((x[-n:], x[:-n]))
<ECODE> Test like: <SCODE>x = torch.arange(5)
print("Orig:", x)
print("Roll 2:", roll(x, 2))
print("Roll -2:", roll(x, -2))
<ECODE> Outputs: <SCODE>Orig: tensor([0, 1, 2, 3, 4])
Roll 2: tensor([3, 4, 0, 1, 2])
Roll -2: tensor([2, 3, 4, 0, 1])
<ECODE> To roll around second axis, use: <SCODE>def roll_1(x, n):  
    return torch.cat((x[:, -n:], x[:, :-n]), dim=1)
<ECODE> It probably can be generalised, but I didn’t need it.

Alexander_Reshytko: <SCODE>def roll(x: torch.Tensor, shift: int, dim: int = -1, fill_pad: Optional[int] = None):

    if 0 == shift:
        return x

    elif shift < 0:
        shift = -shift
        gap = x.index_select(dim, torch.arange(shift))
        if fill_pad is not None:
            gap = fill_pad * torch.ones_like(gap, device=x.device)
        return torch.cat([x.index_select(dim, torch.arange(shift, x.size(dim))), gap], dim=dim)

    else:
        shift = x.size(dim) - shift
        gap = x.index_select(dim, torch.arange(shift, x.size(dim)))
        if fill_pad is not None:
            gap = fill_pad * torch.ones_like(gap, device=x.device)
        return torch.cat([gap, x.index_select(dim, torch.arange(shift))], dim=dim)
<ECODE>

Zuanazzi: I tried to use yours, but I get a compilation error saying that Optional is not defined.

Alexander_Reshytko: <SCODE>from typing import Optional
<ECODE>

tsnowak: <SCODE>from typing import Optional

def roll(x: torch.Tensor, shift: int, dim: int = -1, fill_pad: Optional[int] = None):

    device = x.device
    
    if 0 == shift:
        return x

    elif shift < 0:
        shift = -shift
        gap = x.index_select(dim, torch.arange(shift, device=device))
        if fill_pad is not None:
            gap = fill_pad * torch.ones_like(gap, device=device)
        return torch.cat([x.index_select(dim, torch.arange(shift, x.size(dim), device=device)), gap], dim=dim)

    else:
        shift = x.size(dim) - shift
        gap = x.index_select(dim, torch.arange(shift, x.size(dim), device=device))
        if fill_pad is not None:
            gap = fill_pad * torch.ones_like(gap, device=device)
        return torch.cat([gap, x.index_select(dim, torch.arange(shift, device=device))], dim=dim) 
<ECODE>

lucasb-eyer: