API接口

class roll():

功能描述

使用NPU亲和写法替换swin-transformer中的原生roll。

参数说明

input1 (Tensor) - 输入张量。
shifts (Int或Tuple of python:ints) - 张量元素移动的位置数。如果该shift组成的是tuple，则dims必须是大小相同的tuple，每个维度都将按相应的值roll。
dims (Int或Tuple of python:ints) - 沿着roll的轴。

输出说明

Tensor - shifted input。

示例

调用方式示例：

from torch_npu.contrib.function import roll
shifted_x_npu = roll(input1, shifts=(-shift_size, -shift_size), dims=(1, 2))

使用示例：

       
        >>> input1 = torch.randn(32, 56, 56, 16).npu()
        >>> shift_size = 3
        >>> shifted_x_npu = roll(input1, shifts=(-shift_size, -shift_size), dims=(1, 2))