Repository: pytorch/functorch (language: python)

Issue:
[vmap] Most used `torch.<x>` operations: 150-200
- [x] quantize_per_tensor | 1392
- [ ] diag_embed | 1364 (needs functionalization)
- [x] lerp | 1363
- [x] asin | 1345
- [ ] eig | 1333 (actionable)
- [x] trunc | 1290
- [x] diagonal | 1287
- [x] cosh | 1279
- [x] rfft | 1269
- [x] cumprod | 1260
- [x] addr | 1211
- [x] roll | 1198
- [x] narrow | 1188
- [x] digamma | 1172
- [x] square | 1163
- [x] sinh | 1131
- [x] logspace | 1084
- [x] broadcast_tensors | 1070
- [x] irfft | 1013
- [x] frac | 997
- [x] hann_window | 994
- [x] solve | 989
- [x] logdet | 977
- [x] expm1 | 968
- [x] cdist | 946
- [x] addmv | 903
- [ ] randint_like | 888
- [x] tensordot | 888
- [x] ifft | 877
- [x] true_divide | 854
- [x] erfinv | 830
- [x] addcdiv | 819
- [x] addbmm | 813
- [ ] renorm | 781
- [x] pinverse | 753
- [ ] isclose | 740
- [x] erfc | 729
- [x] is_storage | 725
- [x] triangular_solve | 723
- [x] rot90 | 709
- [x] logical_not | 686
- [x] geqrf | 681
- [x] slogdet | 677
- [ ] lu | 665 (actionable)
- [x] hamming_window | 659
- [ ] orgqr | 651  (actionable)
- [ ] ormqr | 622  (actionable)
- [x] is_floating_point | 602
- [ ] diagflat | 562 (needs OpInfo)
- [ ] cholesky_solve | 559 (actionable)


Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesViews.cpp b/functorch/csrc/BatchRulesViews.cpp
index 7fb2bdbbc..cef2caadd 100644
--- a/functorch/csrc/BatchRulesViews.cpp
+++ b/functorch/csrc/BatchRulesViews.cpp
@@ -148,7 +148,7 @@ std::tuple<Tensor,optional<int64_t>> flip_batch_rule(const Tensor& self, optiona
   auto self_ = moveBatchDimToFront(self, self_bdim);
   VmapDimVector new_dims;
   for (auto i: dims) {
-    new_dims.push_back(getPhysicalDim(self, true, i));
+    new_dims.push_back(getPhysicalDim(self_, true, i));
   }
   return std::make_tuple(at::flip(self_, new_dims), 0);
 }
@@ -283,6 +283,25 @@ std::tuple<Tensor, optional<int64_t>> _reshape_alias_batch_rule(const Tensor& se
   return std::make_tuple(at::reshape(self, new_shape), bdim);
 }
 
+std::tuple<Tensor, optional<int64_t>> roll_batch_rule(const Tensor& self, optional<int64_t> bdim, c10::ArrayRef<long> shifts, c10::ArrayRef<long> dims) {
+  TORCH_INTERNAL_ASSERT(bdim.has_value());
+
+  auto self_ = moveBatchDimToFront(self, bdim);
+  VmapDimVector new_dims;
+  if (!dims.empty()) {
+    for (auto i: dims) {
+      new_dims.push_back(getPhysicalDim(self, true, i));
+    }
+    return std::make_tuple(at::roll(self_, shifts, new_dims), 0);
+  }
+  // We will do something like: t.reshape(a, -1).roll(1, dims=[1, ]).reshape(old_shape)
+  auto old_shape = self_.sizes();
+  new_dims.push_back(1);
+  auto output = at::roll(self_.flatten(1), shifts, new_dims);
+  output = output.reshape(old_shape);
+  return std::make_tuple(output, 0);
+}
+
 std::tuple<Tensor,optional<int64_t>> diagonal_backward_batch_rule(
     const Tensor& grad_input, optional<int64_t> grad_input_bdim,
     IntArrayRef input_sizes, int64_t offset, int64_t dim1, int64_t dim2) {
@@ -339,6 +358,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("squeeze", squeeze_batch_rule);
   VMAP_SUPPORT("squeeze.dim", squeeze_dim_batch_rule);
   VMAP_SUPPORT("_reshape_alias", _reshape_alias_batch_rule);
+  VMAP_SUPPORT("roll", roll_batch_rule);
   VMAP_SUPPORT("diagonal_backward", diagonal_backward_batch_rule);
   VMAP_SUPPORT("select_backward", select_backward_batch_rule);
   VMAP_SUPPORT("slice_backward", slice_backward_batch_rule);


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index abf8be1bb..d625701df 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -452,7 +452,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('put'),
         xfail('quantile'),
         xfail('renorm'),
-        xfail('roll'),
         xfail('scatter_add'),
         xfail('solve'),
         xfail('sort'),
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 434f15d77..97bfe8d22 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2102,6 +2102,15 @@ def test_select(self):
         test(vmap(lambda t: op(t, 1, 1)), (torch.rand(B1, 2, B0, 5),), in_dims=2)
         test(vmap(vmap(lambda t: op(t, 1, 1), in_dims=1)), (torch.rand(B1, 2, B0, B2, 5),), in_dims=2)
 
+    def test_roll_no_dims(self):
+        op = torch.roll
+        test = self._vmap_test
+        B0, B1, B2 = 7, 11, 13
+        test(op, (torch.rand(B0, 2, 5), 2), in_dims=(0, None))
+        test(op, (torch.rand(2, B0, 5), 3), in_dims=(1, None))
+        test(vmap(lambda t: op(t, 3)), (torch.rand(B1, 2, B0, 5),), in_dims=2)
+        test(vmap(vmap(lambda t: op(t, 3), in_dims=1)), (torch.rand(B1, 2, B0, B2, 5),), in_dims=2)
+
     def test_stack(self):
         test = self._vmap_test
         B0, B1 = 5, 7
@@ -3039,7 +3048,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('resize_as_'),
         xfail('resolve_conj'),
         xfail('resolve_neg'),
-        xfail('roll'),
         xfail('scatter'),
         xfail('scatter_add'),
         xfail('take'),