[functorch] Added im2col batch rule and enabled vmap for nn.functiona…

…l.unfold op (pytorch/functorch#262) Description: - Added im2col batch rule and enabled vmap for nn.functional.unfold op - Updated tests Using EXISTING_BDIM macro to put bdim into 0 as im2col expects dim=0 to be batch dim Related to pytorch/functorch#240
zou3519 · Jul 20, 2022 · 13c4548 · 13c4548
1 parent ed20f40
commit 13c4548
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/functorch/functorch/csrc/BatchRulesDecompositions.cpp b/functorch/functorch/csrc/BatchRulesDecompositions.cpp
@@ -87,6 +87,9 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE2(multiply, Tensor );
   OP_DECOMPOSE(narrow);
   OP_DECOMPOSE(negative);
+  OP_DECOMPOSE(nll_loss_nd);
+  OP_DECOMPOSE(nll_loss);
+  OP_DECOMPOSE(nll_loss2d);
   OP_DECOMPOSE2(not_equal, Tensor );
   OP_DECOMPOSE(outer);
   OP_DECOMPOSE(pairwise_distance);
@@ -125,9 +128,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE(var_mean);
   OP_DECOMPOSE2(var_mean, dim);
   OP_DECOMPOSE2(where, self);
-  OP_DECOMPOSE(nll_loss_nd);
-  OP_DECOMPOSE(nll_loss);
-  OP_DECOMPOSE(nll_loss2d);
 }
 
 }}

diff --git a/functorch/functorch/csrc/BatchRulesModules.cpp b/functorch/functorch/csrc/BatchRulesModules.cpp
@@ -396,6 +396,9 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   m.impl("conv2d", convNd_decomp);
   m.impl("conv3d", convNd_decomp);
 
+  EXISTING_BDIM(im2col);
+  EXISTING_BDIM(im2col_backward);
+
   VMAP_SUPPORT("grid_sampler_2d", GRID_SAMPLE_BATCH_RULE(grid_sampler));
   VMAP_SUPPORT("grid_sampler_3d", GRID_SAMPLE_BATCH_RULE(grid_sampler));
   VMAP_SUPPORT("cudnn_grid_sampler", GRID_SAMPLE_BATCH_RULE(cudnn_grid_sampler));

diff --git a/functorch/test/test_ops.py b/functorch/test/test_ops.py
@@ -456,7 +456,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('nn.functional.gelu'),
         xfail('nn.functional.grid_sample'),
         xfail('nn.functional.pad', 'circular'),
-        xfail('nn.functional.unfold'),
         xfail('norm', 'fro'),
         xfail('norm', 'inf'),
         xfail('norm', 'nuc'),

diff --git a/functorch/test/test_vmap.py b/functorch/test/test_vmap.py
@@ -3124,7 +3124,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('nn.functional.cross_entropy', 'none'),
         xfail('nn.functional.cross_entropy', 'sum'),
         xfail('nn.functional.pad', 'circular'),
-        xfail('nn.functional.unfold'),
         xfail('norm', 'fro'),
         xfail('norm', 'nuc'),
         xfail('ormqr'),