Roll operator t32802531 #13261

nairbv · 2018-10-29T19:00:12Z

Adding a roll operator

aten/src/ATen/native/TensorTransformations.cpp

aten/src/ATen/native/cuda/TensorTransformations.cu

aten/src/ATen/native/TensorTransformations.cpp

+  // If the first dimension is zero, this is an empty tensor and rolls do nothing.
+  // Return a clone so the caller can safely modify result, and avoid a div by
+  // zero error below.
+  if( self.size(0) == 0 ) {


aten/src/ATen/native/TensorTransformations.cpp

+Tensor roll_cpu(const Tensor& self, IntList shifts, IntList dims) {
+  // todo: support rolling along no or multiple dimensions as in numpy.roll.
+  AT_CHECK(dims.size() == 1, "only single dimension roll currently supported");
+  AT_CHECK(shifts.size() == dims.size(), "shifts and dimensions must align");


aten/src/ATen/native/TensorTransformations.cpp

+    vec[index++] = tensors[i];
+  }
+
+  auto stacked = at::stack(vec, dim);


aten/src/ATen/native/native_functions.yaml

@@ -1672,6 +1672,12 @@
    CPU: flip_cpu
    CUDA: flip_cuda

+- func: roll(Tensor self, IntList[1] shifts, IntList[1] dims) -> Tensor


aten/src/ATen/native/cuda/TensorTransformations.cu

+  // corrects the difference.
+  if( start < 0 ) start = start + size;
+
+  const int64_t block_size = 512;


aten/src/ATen/native/cuda/TensorTransformations.cu

+    return;
+  }
+  // roll dim idx is the index of linear_index along the rolling dimension.
+  int64_t roll_dim_idx = linear_index % (stride * size) / stride;


torch/_torch_docs.py

@@ -4620,6 +4620,39 @@ def parse_kwargs(desc):
             [ 0,  1]]])
 """)

+add_docstr(torch.roll,
+           r"""
+roll(input, shift, dims) -> Tensor


torch/_tensor_docs.py

@@ -888,6 +888,13 @@ def add_docstr_all(method, docstr):
 See :func:`torch.flip`
 """)

+add_docstr_all('roll',
+               r"""
+roll(shift, dims) -> Tensor


aten/src/ATen/native/cuda/TensorTransformations.cu

+    return;
+  }
+  // roll dim idx is the index of linear_index along the rolling dimension.
+  int64_t roll_dim_idx = linear_index % (stride * size) / stride;


gchanan

I think this is good to go once the shift/shifts issue is dealt with.

facebook-github-bot

@nairbv is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

Summary: Adding a roll operator Pull Request resolved: pytorch/pytorch#13261 Differential Revision: D12922575 Pulled By: nairbv fbshipit-source-id: ff05c075d9c484a615011192b023debf47da4017

nairbv added 4 commits October 25, 2018 11:45

add cpu version of roll operator.

8a64846

add cuda version of roll operator

d062069

fix roll_cuda and add autograd tests

983c364

cleanup for roll operator interfaces

9c47d7a

ssnl reviewed Oct 29, 2018

View reviewed changes

aten/src/ATen/native/TensorTransformations.cpp Outdated Show resolved Hide resolved

ssnl reviewed Oct 29, 2018

View reviewed changes

aten/src/ATen/native/cuda/TensorTransformations.cu Outdated Show resolved Hide resolved

nairbv added 4 commits October 29, 2018 12:22

Merge branch 'master' into roll_operator_T32802531

9840a78

resolve a merge conflict

29e20ba

improve performance of roll_cpu using unbind instead of select

69a4608

fix lint errors

131d9cc

gchanan reviewed Oct 29, 2018

View reviewed changes

aten/src/ATen/native/cuda/TensorTransformations.cu Outdated Show resolved Hide resolved

nairbv added 3 commits October 30, 2018 13:36

fix roll operator for strided tensors

df2c956

Merge branch 'roll_operator_T32802531' into roll_operator_T32802531

5730d0c

remove unused code

7270aaa

gchanan reviewed Oct 31, 2018

View reviewed changes

nairbv added 4 commits November 1, 2018 14:08

some updates to roll operator from code review

f3eb196

simplify return statement

623bf65

use numel instead of size

7b437d6

use arange instead of fully specifying tensor

3ab3061

gchanan reviewed Nov 2, 2018

View reviewed changes

gchanan approved these changes Nov 2, 2018

View reviewed changes

fix docs for roll

af59b01

facebook-github-bot reviewed Nov 5, 2018

View reviewed changes

facebook-github-bot closed this in 07f8b61 Nov 5, 2018

zou3519 mentioned this pull request Jan 7, 2019

[feature request] torch.roll like numpy.roll #10434

Closed

ezyang added the merged label Jun 25, 2019

Roll operator t32802531 #13261

Roll operator t32802531 #13261

Uh oh!

Conversation

nairbv commented Oct 29, 2018

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

gchanan left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!