Support vmapvjp through F.pad (#357)

Author: zou3519

functorch/csrc/BatchingRegistrations.cpp

@@ -368,6 +368,10 @@ static void checkBatchDimsAtFrontInLayout(IntArrayRef physical_strides, int64_t
     // No example dimensions
     return;
   }
+  if (num_batch_dims == 1 && physical_strides.size() > 0 && physical_strides[0] == 0) {
+    // degenerate batch dim
+    return;
+  }
   TORCH_CHECK(*smallest_batch_stride >= *largest_example_stride,
     "vmap: Calling Tensor.as_strided is not supported unless the batch dims being ",
     "vmapped over are at the front of the tensor (in memory layout). When they are ",
@@ -929,7 +933,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
 // //   // Tensor.new_* operators
 //   m.impl("ones_like", ones_like_batching_rule);
 // //   m.impl("new_empty", new_empty_batching_rule);
-//   m.impl("new_empty_strided", new_empty_strided_batching_rule);
+  m.impl("new_empty_strided", new_empty_strided_batching_rule);
 // //   m.impl("new_zeros", new_zeros_batching_rule);
 // //
   m.impl("contiguous", contiguous_batching_rule);

test/test_ops.py

@@ -623,9 +623,6 @@ def test_vmapvjp(self, device, dtype, op):
         # xfail "above the line".
         xfail('double', 'channels_last'),
 
-        # See https://github.com/pytorch/pytorch/issues/66357
-        xfail('nn.functional.pad', 'circular'),
-
         # RuntimeError: expand: the number of sizes provided (1) must be greater or
         # equal to the number of dimensions in the tensor (2)
         xfail('nanquantile'),
@@ -690,7 +687,6 @@ def test_vmapjvp(self, device, dtype, op):
         xfail('linalg.inv'),
         xfail('masked_fill'),
         xfail('linalg.tensorinv'),
-        xfail('nn.functional.pad', 'circular'),
         xfail('linalg.matrix_power'),
         xfail('maximum'),
         xfail('linalg.householder_product'),
@@ -796,7 +792,6 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('nanquantile'),
         xfail('nn.functional.conv_transpose2d'),
         xfail('nn.functional.gelu'),
-        xfail('nn.functional.pad', 'circular'),
         xfail('norm', 'nuc'),
         xfail('pinverse'),
         xfail('prod'),

test/test_vmap.py

@@ -2177,8 +2177,6 @@ def test_new_empty(self):
         result = vmap(vmap(lambda x: op(x, [2, 3])))(torch.randn(B0, B1))
         self.assertEqual(result.shape, [B0, B1, 2, 3])
 
-    # TODO: new_empty_strided BR
-    @unittest.expectedFailure
     def test_new_empty_strided(self):
         # Empty is non-deterministic so we just check that the size and shape
         # of the output are what we expect and that the vmap fallback isn't used