Implementation of function like `numpy.roll`

username: Ja-Keoung_Koo
cooked: <p>I am trying to implement a function similar to <code>numpy.roll</code> function. To use cuda, I want to use only torch Tensors, but it seems to be hard.</p>
<p>Ultimately, I want to implement image gradient with forward difference and Neumann boundary condition. For example, the numpy version of it is as follows:</p>
<pre><code class="lang-auto">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
</code></pre>
<p>I tried to use <code>[-1, 1]</code> filter, using <code>torch.nn.Functional.conv2d</code>. Because there is no boundary option except zero padding and the filter is even filter, it looks complicated.</p>
<p>Is there someone who can help implement one like <code>np.roll</code>?</p>

username: Ja-Keoung_Koo
cooked: <p>Thanks for your reply, although I checked before. As you said, it seems hard to implement.</p>

username: apaszke
cooked: <p>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.</p>

username: Ryan_Burgert
cooked: <p>i made an account to say that i think thats a bit yucky, it’s nicer to read roll <img src="//discuss.pytorch.org/images/emoji/apple/frowning.png?v=5" title=":frowning:" class="emoji" alt=":frowning:"></p>

username: azban
cooked: <p>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</p>
<pre><code class="lang-python">def roll(tensor, shift, axis):
    if shift == 0:
        return tensor

    if axis &lt; 0:
        axis += tensor.dim()

    dim_size = tensor.size(axis)
    after_start = dim_size - shift
    if shift &lt; 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)
</code></pre>

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

username: Alexander_Reshytko
cooked: <p><a class="mention" href="/u/jaromiru">@jaromiru</a> thank you!<br>
The solution below is a generalization of yours to an arbitrary axis:</p>
<pre><code class="lang-auto">def roll(x: torch.Tensor, shift: int, dim: int = -1, fill_pad: Optional[int] = None):

    if 0 == shift:
        return x

    elif shift &lt; 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)
</code></pre>

username: Zuanazzi
cooked: <p>I tried to use yours, but I get a compilation error saying that Optional is not defined.</p>

username: Alexander_Reshytko
cooked: <p><a class="mention" href="/u/zuanazzi">@Zuanazzi</a>,</p>
<pre><code class="lang-python">from typing import Optional
</code></pre>

username: tsnowak
cooked: <p>Extending the solution to support devices <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<pre><code class="lang-auto">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 &lt; 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) 
</code></pre>

username: lucasb-eyer
cooked: <p>Since this is still getting answers in 2020 and is the top Google answer, it’s worth pointing out that there is now a proper <a href="https://pytorch.org/docs/stable/torch.html#torch.roll" rel="nofollow noopener"><code>torch.roll</code></a> function.</p>