Random assortment of batch rules and cleanups (#347)

Test Plan:
- wait for tests

Author: zou3519

functorch/csrc/BatchRulesBinaryOps.cpp

@@ -334,7 +334,9 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
 
 #define LOGICAL_COMPARISON_POINTWISE(op) \
   VMAP_SUPPORT(#op, \
-      SINGLE_ARG(comparison_pointwise_batch_rule<decltype(&ATEN_FN(op)), &ATEN_FN(op)>));
+      SINGLE_ARG(comparison_pointwise_batch_rule<decltype(&ATEN_FN(op)), &ATEN_FN(op)>)); \
+  m.impl(#op "_", inplacePlumbing2< \
+     DECLTYPE_AUTO(&binary_pointwise_inplace_batch_rule<TensorInplaceT, &Tensor:: op ## _ >)>);
 
   LOGICAL_COMPARISON_POINTWISE(logical_and);
   LOGICAL_COMPARISON_POINTWISE(logical_or);

functorch/csrc/BatchRulesDecompositions.cpp

@@ -227,6 +227,14 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE(conv_transpose1d);
   OP_DECOMPOSE2(conv_transpose2d, input);
   OP_DECOMPOSE2(conv_transpose3d, input);
+  OP_DECOMPOSE(conv1d);
+  OP_DECOMPOSE(conv2d);
+  OP_DECOMPOSE(conv3d);
+  OP_DECOMPOSE2(conv1d, padding);
+  OP_DECOMPOSE2(conv2d, padding);
+  OP_DECOMPOSE2(conv3d, padding);
+  OP_DECOMPOSE(_convolution_mode);
+  OP_DECOMPOSE(type_as);
   DECOMPOSE_FUNCTIONAL(diag_embed);
   DECOMPOSE_FUNCTIONAL(block_diag);
 }

functorch/csrc/BatchRulesModules.cpp

@@ -87,10 +87,20 @@ convolution_batch_rule(const Tensor& lhs, optional<int64_t> lhs_bdim, const Tens
     return result;
   }
 }
-Tensor convNd_decomp(const Tensor &self, const Tensor &weight, const optional<Tensor>& bias, IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
-  std::vector<int64_t> t(self.dim() - 2, 0);
-  IntArrayRef out_padding(t);
-  return at::convolution(self, weight, bias, stride, padding, dilation, false, out_padding, groups);
+
+Tensor _convolution_decomp(
+    const Tensor& input_r, const Tensor& weight_r, const c10::optional<Tensor>& bias_r_opt,
+    IntArrayRef stride_, IntArrayRef padding_, IntArrayRef dilation_,
+    bool transposed_, IntArrayRef output_padding_, int64_t groups_,
+    bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) {
+  // Ignore everything. If the user called this in the normal way,
+  // then they should be fine.
+  (void*) benchmark;
+  (void*) deterministic;
+  (void*) cudnn_enabled;
+  (void*) allow_tf32;
+  return at::convolution(
+      input_r, weight_r, bias_r_opt, stride_, padding_, dilation_, transposed_, output_padding_, groups_);
 }
 
 // Tensor convNd_transpose_decomp(const Tensor &self, const Tensor &weight, const optional<Tensor>& bias, IntArrayRef stride, IntArrayRef padding, IntArrayRef dilation, int64_t groups) {
@@ -517,13 +527,10 @@ struct CudnnGridSampleBackwardBatchRuleHelper {
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("convolution", convolution_batch_rule);
-  // m.impl("conv_transpose2d", convNd_transpose_decomp);
+  m.impl("_convolution", _convolution_decomp);
   m.impl("mkldnn_convolution", mkldnn_convolution_decomp);
   m.impl("cudnn_convolution_backward", cudnn_convolution_backward_plumbing);
   m.impl("cudnn_convolution", cudnn_convolution_plumbing);
-  m.impl("conv1d", convNd_decomp);
-  m.impl("conv2d", convNd_decomp);
-  m.impl("conv3d", convNd_decomp);
 
   EXISTING_BDIM(im2col);
   EXISTING_BDIM(im2col_backward);

test/discover_coverage.py

@@ -359,9 +359,13 @@ def print_coverage_info(th=100, nn=25):
     statuses = transpose_statuses(get_top_ops(th, nn), invert=True)
     top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(th, nn)
 
+    # testing problems
+    exemptions = {
+        'torch.nn.functional.dropout', # randomness
+    }
+
     # Allowed exemptions
     vmap_exemptions = {
-        'torch.nn.functional.dropout', # randomness
         'torch.randn_like', # randomness
         'torch.allclose', # number output
         'torch.unique', # dynamic
@@ -374,6 +378,8 @@ def print_coverage_info(th=100, nn=25):
     remove_from_set(statuses['test_vmapvjp_has_batch_rule'], vmap_exemptions)
     remove_from_set(statuses['test_op_has_batch_rule'], vmap_exemptions)
     remove_from_set(statuses['test_vmapjvp'], vmap_exemptions)
+    for test in tests:
+        remove_from_set(statuses[test], exemptions)
 
     print(f"total ops in set: {th + nn}")
     print(f"tested by OpInfo: {th + nn - len(top_ops_not_covered_by_opinfo)}")

test/test_ops.py

@@ -475,7 +475,6 @@ def vjp_of_vjp(*args_and_cotangents):
         xfail('double', 'channels_last'),
         xfail('nn.functional.gaussian_nll_loss'),
         xfail('nn.functional.poisson_nll_loss'),
-        skip('nn.functional.conv1d', device_type='cuda'),
         xfail('fft.rfft2'),
         xfail('lu'),
         skip('qr'),  # Nondetermistic
@@ -694,10 +693,8 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('__getitem__'),
         xfail('cdist'),
         xfail('cholesky'),
-        xfail('clamp', 'scalar'),
         xfail('complex'),
         xfail('copysign'),
-        xfail('corrcoef'),
         xfail('cummax'),
         xfail('cummin'),
         xfail('cumprod'),
@@ -732,7 +729,6 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('linalg.slogdet'),
         xfail('linalg.solve'),
         xfail('linalg.tensorinv'),
-        xfail('linalg.vector_norm'),
         xfail('logdet'),
         xfail('lu'),
         xfail('lu_solve'),
@@ -741,16 +737,11 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('masked_scatter'),
         xfail('masked_select'),
         xfail('matrix_exp'),
-        xfail('max', 'reduction_no_dim'),
-        xfail('median'),
-        xfail('min', 'reduction_no_dim'),
-        xfail('nanmedian'),
         xfail('nanquantile'),
         xfail('nn.functional.conv_transpose2d'),
         xfail('nn.functional.gelu'),
         xfail('nn.functional.pad', 'circular'),
         xfail('norm', 'fro'),
-        xfail('norm', 'inf'),
         xfail('norm', 'nuc'),
         xfail('pinverse'),
         xfail('prod'),
@@ -785,8 +776,6 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('nn.functional.instance_norm'),
         xfail('nn.functional.poisson_nll_loss'),
         xfail('nn.functional.conv_transpose3d'),
-        xfail('_masked.norm'),
-        xfail('_masked.normalize'),
         xfail('nn.functional.bilinear'),
         xfail('nn.functional.prelu'),
         xfail('nn.functional.glu'),

test/test_vmap.py

@@ -3221,7 +3221,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('nanmean'),
         xfail('vstack'),
         xfail('nn.functional.dropout'),
-        xfail('nn.functional.conv2d', ''),
         xfail('nn.functional.batch_norm'),
         xfail('resize_'),
         xfail('view_as_complex'),
@@ -3253,7 +3252,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('short', 'channels_last'),
         xfail('unique_consecutive'),
         xfail('unique'),
-        xfail('nn.functional.conv1d'),
         xfail('nn.functional.cosine_embedding_loss'),
         xfail('nn.functional.ctc_loss'),
         xfail('nn.functional.gaussian_nll_loss'),