Repository: pytorch/functorch (language: python)

Issue:
Actually actionable list of batching rules to write
For each of the items here, we should make sure all compositions (vmap, vmap x vjp) have a batching rule. All of these items should be actionable (in that it is possible to write a batching rule and we are not blocked on functionalization, which is coming soon).

Note: you may need to write an OpInfo for the operator if it doesn't exist already or wait for one to be added. A lot of folks are adding OpInfos right now, so if the OpInfo doesn't exist please ask first to see if someone is working on it.

Note: if any of the operations decompose into in-place operations, then we need functionalization to handle them. I think I've already filtered out all of those, but please check me on that.

Parcel 1: top nn.functional.* and top torch.* foo
- [x] torch.nn.functional.interpolate (this involves adding batching rules for adaptive_avg_pool1d, but we might as well do `adaptive_avg_pool{1, 2, 3}d` as well as their backward variants while we're at it)
- [x] nn.functional.unfold. Involves writing a batching rule for im2col im2col_backward (https://github.com/pytorch/functorch/pull/262)
- [x] nn.functional.grid_sample. Involves writing a batching rule for the backward operator. @vfdev-5 
- [x] torch.pow. The backward needs batching rule for logical_and ; use this as an opportunity to write the batching rules for `logical_{and, or, xor}` if those don't exist yet. We may need to also add a change to PyTorch core to make the logical_* functions primitives w.r.t. autograd

Parcel 2: new_blah
- [x] new_empty, new_full, new_ones, new_zeros, empty_like, zeros_like, ones_like, full_like
- [x] `adaptive_max_pool{1, 2, 3}d` as well as the backward variants
- [x] diagonal_scatter, select_scatter, slice_scatter @kshitij12345 
- [ ] linalg.householder_product (forward pass only)
- [x] pixel_shuffle, pixel_unshuffle (these might just be OP_DECOMPOSE)
- [x] isinf, isfinite, isnan
- [x] _cdist_forward, _cdist_backward (try to write a batching rule if possible. If not possible, we may need to write a decomposition) @vfdev-5 

Parcel 3: linalg things
- [x] _lu_with_info (forward pass only) @vfdev-5 
- [x] linalg.eig (forward pass only) @vfdev-5 
- [x] torch.addr @vfdev-5 
- [x] cholesky_solve (forward pass only) @vfdev-5 

Parcel 4:
- [x] ~~index_select, index_copy, etc, all need a backward formula in pytorch/pytorch https://github.com/pytorch/functorch/issues/260~~


Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesDecompositions.cpp b/functorch/csrc/BatchRulesDecompositions.cpp
index 1528e914e..eb230315e 100644
--- a/functorch/csrc/BatchRulesDecompositions.cpp
+++ b/functorch/csrc/BatchRulesDecompositions.cpp
@@ -18,7 +18,10 @@ namespace at { namespace functorch {
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE(absolute);
+  OP_DECOMPOSE(avg_pool1d);
+  OP_DECOMPOSE(adaptive_avg_pool1d);
   OP_DECOMPOSE(adaptive_avg_pool2d);
+  OP_DECOMPOSE(adaptive_avg_pool3d);
   OP_DECOMPOSE(arccos);
   OP_DECOMPOSE(arccosh);
   OP_DECOMPOSE(arcsin);
diff --git a/functorch/csrc/BatchRulesPooling.cpp b/functorch/csrc/BatchRulesPooling.cpp
index 071df32ef..df144bf70 100644
--- a/functorch/csrc/BatchRulesPooling.cpp
+++ b/functorch/csrc/BatchRulesPooling.cpp
@@ -84,8 +84,13 @@ max_pool2d_with_indices_batch_rule(
 
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   EXISTING_BDIM(_adaptive_avg_pool2d);
+  EXISTING_BDIM_ALL_BOXED(_adaptive_avg_pool2d_backward);
+  EXISTING_BDIM(_adaptive_avg_pool3d);
+  EXISTING_BDIM_ALL_BOXED(_adaptive_avg_pool3d_backward);
   EXISTING_BDIM(avg_pool2d);
+  EXISTING_BDIM(avg_pool3d);
   EXISTING_BDIM_ALL_BOXED(avg_pool2d_backward);
+  EXISTING_BDIM_ALL_BOXED(avg_pool3d_backward);
   VMAP_SUPPORT("max_pool2d_with_indices", max_pool2d_with_indices_batch_rule);
   VMAP_SUPPORT("max_pool2d_with_indices_backward", max_pool2d_with_indices_backward_batch_rule);
 }


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index d55603579..1704e062c 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -461,11 +461,9 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('msort'),
         xfail('nanmedian'),
         xfail('nanquantile'),
-        xfail('nn.functional.adaptive_avg_pool2d'),
         xfail('nn.functional.conv_transpose2d'),
         xfail('nn.functional.gelu'),
         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'),
@@ -505,9 +503,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('fft.ihfft2'),
         xfail('fft.ihfftn'),
         xfail('fft.rfft2'),
-        xfail('nn.functional.adaptive_avg_pool1d'),
-        xfail('nn.functional.adaptive_avg_pool3d'),
-        xfail('nn.functional.avg_pool3d'),
         xfail('nn.functional.embedding'),
     }))
     def test_vmapvjp_has_batch_rule(self, device, dtype, op):
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 8cffb0674..5d47d65d4 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -3094,7 +3094,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('nn.functional.cross_entropy', 'mean'),
         xfail('nn.functional.cross_entropy', 'none'),
         xfail('nn.functional.cross_entropy', 'sum'),
-        xfail('nn.functional.interpolate', 'area'),
         xfail('nn.functional.pad', 'circular'),
         xfail('nn.functional.unfold'),
         xfail('norm', 'fro'),
@@ -3143,10 +3142,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('fft.rfft2'),
         xfail('isinf'),
         xfail('isreal'),
-        xfail('nn.functional.adaptive_avg_pool1d'),
-        xfail('nn.functional.adaptive_avg_pool3d'),
-        xfail('nn.functional.avg_pool1d'),
-        xfail('nn.functional.avg_pool3d'),
         xfail('nn.functional.pixel_shuffle'),
         xfail('nn.functional.pixel_unshuffle'),
     }))