Implementation of function like `numpy.roll`
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.
Ja-Keoung_Koo
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.
apaszke
Ryan_Burgert
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>
azban
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.
jaromiru
<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>
Alexander_Reshytko
I tried to use yours, but I get a compilation error saying that Optional is not defined.
Zuanazzi
<SCODE>from typing import Optional
<ECODE>
Alexander_Reshytko
<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>
tsnowak
lucasb-eyer