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:
@zou3519 `nn.functional.pad` with circular option requires to do a copy: `out[..., out_d0:out_d1] = input[..., in_d0:in_d1]` and thus there is the following error is raised:
```
>           out[..., out_d0:out_d1] = input[..., in_d0:in_d1]                                                                                                                               
E           RuntimeError: vmap: aten::copy_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over at level 2. Please try to use out-of-place operators instead of aten::copy_. If said operator is being called inside the PyTorch framework, please file a bug report instead. 
```
Can we do something with that ?
Aha. Nope, we can't do anything about that until functionalization is in, good catch.
Hi @zou3519, the "forward pass only" ops above means that the vjp and related operators require the functionalization too?
> Hi @zou3519, the "forward pass only" ops above means that the vjp and related operators require the functionalization too?

Yes, "forward pass only" means we should only try to get the `vmap` tests passing and none of the vjp/grad/compositions of {vjp, grad} tests.
@kshitij12345 on which tasks from Parcel 2 you are working on and plan to work on ?


I can start working on `_cdist_forward`, `_cdist_backward`
@vfdev-5 I think I'll be picking `diagonal_scatter` next. Go ahead with `_cdist_forward, _cdist_backward`

Fix patch:
diff --git a/functorch/csrc/BatchRulesBinaryOps.cpp b/functorch/csrc/BatchRulesBinaryOps.cpp
index 5de8b5285..3f4894944 100644
--- a/functorch/csrc/BatchRulesBinaryOps.cpp
+++ b/functorch/csrc/BatchRulesBinaryOps.cpp
@@ -155,6 +155,14 @@ std::tuple<Tensor,optional<int64_t>> masked_select_batch_rule(
   return std::make_tuple(result, 0);
 }
 
+Tensor addr_decomposition(
+    const Tensor& self, const Tensor& vec1, const Tensor& vec2,
+    const Scalar& beta, const Scalar& alpha) {
+
+  auto outer = alpha * vec1.unsqueeze(-1) * vec2.unsqueeze(-2);
+  return self * beta + outer;
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
 #define BINARY_POINTWISE2(op, overload) \
   VMAP_SUPPORT(#op"."#overload, BINARY_POINTWISE_BATCH_RULE(ATEN_FN2(op, overload)));
@@ -193,6 +201,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   BINARY_SCALAR_2(add, Tensor, Scalar);
   POINTWISE_BOXED(addcdiv);
   POINTWISE_BOXED(addcmul);
+  m.impl("addr", addr_decomposition);
   BINARY_POINTWISE(atan2);
   BINARY_SCALAR_2(bitwise_and, Tensor, Scalar);
   BINARY_POINTWISE2(bitwise_or, Tensor);


Test patch:
diff --git a/test/test_ops.py b/test/test_ops.py
index 29cfdd9ad..54b74c093 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -570,7 +570,6 @@ def test_vmapjvp(self, device, dtype, op):
     @skipOps('TestOperators', 'test_vmapvjp_has_batch_rule', vmapvjp_fail.union({
         xfail('view_as_complex'),
         xfail('__getitem__'),
-        xfail('addr'),
         xfail('cdist'),
         xfail('cholesky'),
         xfail('clamp'),
diff --git a/test/test_vmap.py b/test/test_vmap.py
index db97de351..2603e69d9 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -3099,7 +3099,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
 
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestVmapOperatorsOpInfo', 'test_op_has_batch_rule', vmap_fail.union({
-        xfail('addr'),
         xfail('cdist'),
         xfail('complex'),
         xfail('copysign'),