port and fix unfold batching rule (#206)

Author: kshitij12345

functorch/csrc/BatchRulesScatterOps.cpp

@@ -293,7 +293,6 @@ std::tuple<Tensor,optional<int64_t>> scatter_reduce_batch_rule(
     const Tensor& index, optional<int64_t> index_bdim,
     const Tensor& src, optional<int64_t> src_bdim,
     const c10::string_view reduce) {
-  using scatter_reduce_value_sig = Tensor (*)(const Tensor&, int64_t, const Tensor&, const Tensor&, const c10::string_view reduce);
   return scatter_batch_rule(ATEN_FN2(scatter, reduce),
                             self, self_bdim, dim, index, index_bdim, src, src_bdim, reduce);
 }
@@ -304,7 +303,6 @@ std::tuple<Tensor,optional<int64_t>> scatter_value_reduce_batch_rule(
     const Tensor& index, optional<int64_t> index_bdim,
     const Scalar& src,
     const c10::string_view reduce) {
-  using scatter_reduce_value_sig = Tensor (*)(const Tensor&, int64_t, const Tensor&, const Scalar&, const c10::string_view reduce);
   return scatter_batch_rule(ATEN_FN2(scatter, value_reduce),
                             self, self_bdim, dim, index, index_bdim, src, reduce);
 }

functorch/csrc/BatchRulesViews.cpp

@@ -429,6 +429,23 @@ std::tuple<Tensor, optional<int64_t>> expand_batch_rule(
   return std::make_tuple(self_.view(view_shape).expand(size_, implicit), 0);
 }
 
+std::tuple<Tensor, optional<int64_t>> unfold_batch_rule(
+    const Tensor &self, optional<int64_t> self_bdim, int64_t dim, int64_t size, int64_t step)
+{
+  TORCH_INTERNAL_ASSERT(self_bdim.has_value());
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto logical_rank = rankWithoutBatchDim(self, self_bdim);
+  dim = maybe_wrap_dim(dim, logical_rank) + 1;
+  if (logical_rank==0) {
+    self_ = self_.unsqueeze(-1);
+  }
+  auto result = self_.unfold(dim, size, step);
+  if (logical_rank==0) {
+    result = result.squeeze(-1);
+  }
+  return std::make_tuple(result, 0);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("diag", diag_batch_rule);
   VMAP_SUPPORT("chunk", chunk_batching_rule);
@@ -453,6 +470,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("slice_backward", slice_backward_batch_rule);
   VMAP_SUPPORT("view", view_batching_rule);
   VMAP_SUPPORT("expand", expand_batch_rule);
+  VMAP_SUPPORT("unfold", unfold_batch_rule);
 }
 
 }}

functorch/csrc/BatchingRegistrations.cpp

@@ -356,17 +356,6 @@ std::vector<Tensor> unbind_batching_rule(const Tensor& self, int64_t dim) {
   return result;
 }
 
-Tensor unfold_batching_rule(const Tensor& self, int64_t dim, int64_t size, int64_t step) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.unfold(dim, size, step);
-  }
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto dim_physical = self_physical.getPhysicalDim(dim);
-  auto result = self_physical.tensor().unfold(dim_physical, size, step);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
 Tensor contiguous_batching_rule(const Tensor& self, MemoryFormat memory_format) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
@@ -909,7 +898,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   // m.impl("trace", trace_batching_rule);
   m.impl("transpose.int", transpose_int_batching_rule);
   m.impl("unbind.int", unbind_batching_rule);
-  m.impl("unfold", unfold_batching_rule);
   m.impl("unsqueeze_", unsqueeze__batching_rule);
   m.impl("view_as", native::view_as); // composite wrt autograd
 

test/test_ops.py

@@ -355,7 +355,6 @@ def vjp_of_vjp(*args_and_cotangents):
         xfail('quantile'),
         xfail('symeig'),
         xfail('take'),
-        xfail('unfold'),
         xfail('linalg.tensorinv'),
         xfail('nn.functional.conv_transpose2d', device_type='cuda'),
         xfail('nanmean'),
@@ -527,7 +526,6 @@ def test():
         xfail('gradient'),
         xfail('hsplit'),
         xfail('nn.functional.pad', 'circular'),
-        xfail('unfold'),
         xfail('vsplit'),
         xfail('dstack'),
         xfail('hstack'),

test/test_vmap.py

@@ -3005,7 +3005,6 @@ class TestVmapOperatorsOpInfo(TestCase):
 
         # entries in here don't work and need to be fixed.
         # Each one of these is a bug
-        xfail('unfold'),
         xfail('svd', device_type='cuda'),
         xfail('linalg.svd', device_type='cuda'),
         xfail('index_put'),
@@ -3097,7 +3096,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('take_along_dim'),
         xfail('tensor_split'),
         xfail('to_sparse'),
-        xfail('unfold'),
         xfail('vdot'),
         xfail('vsplit'),
         xfail('__getitem__'),