[functorch] Added backward batch rule for pad replicate/reflect modes (…

…pytorch/functorch#251) Description: - Added backward batch rule for pad replicate/reflect modes - Updated tests Related to pytorch/functorch#240
zou3519 · Jul 20, 2022 · 05ec481 · 05ec481
1 parent 7a42ec1
commit 05ec481
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/functorch/functorch/csrc/BatchRulesModules.cpp b/functorch/functorch/csrc/BatchRulesModules.cpp
@@ -409,6 +409,14 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   EXISTING_BDIM(replication_pad2d);
   EXISTING_BDIM(replication_pad3d);
 
+  EXISTING_BDIM_ALL_BOXED(replication_pad1d_backward);
+  EXISTING_BDIM_ALL_BOXED(replication_pad2d_backward);
+  EXISTING_BDIM_ALL_BOXED(replication_pad3d_backward);
+
+  EXISTING_BDIM_ALL_BOXED(reflection_pad1d_backward);
+  EXISTING_BDIM_ALL_BOXED(reflection_pad2d_backward);
+  EXISTING_BDIM_ALL_BOXED(reflection_pad3d_backward);
+
   UPSAMPLE_BATCH(upsample_bicubic2d);
   UPSAMPLE_BATCH(upsample_bilinear2d);
   UPSAMPLE_BATCH(upsample_linear1d);

diff --git a/functorch/test/test_ops.py b/functorch/test/test_ops.py
@@ -465,8 +465,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('nn.functional.grid_sample'),
         xfail('nn.functional.interpolate', 'area'),
         xfail('nn.functional.pad', 'circular'),
-        xfail('nn.functional.pad', 'reflect'),
-        xfail('nn.functional.pad', 'replicate'),
         xfail('nn.functional.unfold'),
         xfail('norm', 'fro'),
         xfail('norm', 'inf'),