Repository: pytorch/functorch (language: python)

Issue:
squeeze.dim batching rule doesn't work for scalar tensors
```
def f(x):
    return torch.squeeze(x, dim=0)

vmap(f)(torch.randn((3)))
```


Discussion / hints:
PyTorch has really weird special cases around scalar tensors. Not sure how easy it will be for us to replicate the semantics.
> Not sure how easy it will be for us to replicate the semantics.

I suspect we're going to need to do so regardless.
@zou3519 - what's the expected behavior in this case? I'm trying to figure out what this should actually do.

Reasoning it out:

1. `x = torch.randn((3))` gives a tensor with shape `(3) (dim=1)`.
2. `f(x) : Tensor<{arbitrary shape; dim >= 1}> -> Tensor<{resultant shape; dim >= 1}>`
3. `vmap(f(x)) : Tensor<{arbitrary shape + 1 dim}> -> Tensor<{resultant shape + 1 dim}>`

The types/shapes don't match up for `x`. E.g. `x` only has a single dimension when passed to `vmap(f(x))` and would thus fail? Or should `x` be promoted to `(3, 1)` automatically (or `(1, 3)`) and then passed to `vmap(f(x))`?

This seems like an issue in the general dispatcher (rather than a batching rule registration style). E.g. when you promote a function using `vmap` the resultant function should check for scalar inputs (i.e. `dim=1`) and automatically promote to `dim=2`?
@laurencer good question. For out-of-place operations and ignoring views, vmap(f)(x) should be equivalent to running `torch.stack([f(xi) for xi in x.unbind(0)])`. This heuristic tells us the following:
1. `x = torch.randn((3))` gives a tensor with shape `(3,)` (dim=1)
2. `f(x[0])` gives a tensor with shape `[]` (dim=0), so `torch.stack([f(xi) for xi in x.unbind(0)])` gives a tensor of shape `[3]` (dim=1)
3. so `vmap(f)(x)` should give us a tensor with shape `(3,)` (dim=1)

I think the batching rule for `squeeze` should check to see if the tensor has dim 1 (and that the dim argument is equal to 0). If it does, then it returns an alias of the tensor (via tensor.alias()).
@zou3519 - is there a description of the new versus old style of batching rules?

I'm having a bit of trouble understanding what the following does and whether it's still needed in the new-style:

```cpp
if (!participatesInCurrentLevel(self)) {
    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
    return self.squeeze();
}
```

From the docs - it seems as though it's now a more functional approach where raw/native Torch tensors are passed in (instead of wrappers) along with optional parameters that describe the batch dimension (instead of being carried on the wrappers).

Is it right to assume that the `optional<uint_64>` parameter passed after every Tensor is the number of batch dimensions (e.g. if it is set to `2` then the first `2` dimensions of the Tensor are batch dimensions?).

If it is empty then there are no batch dimensions?

Also if I return an empty batch dimension - then it appears to assume that there is `1` batch dimension (or this might be the original input number)?

Also there's implementations for `squeeze` and `squeeze.dim` - through trial-and-error I figured out that `squeeze.dim` is invoked when the `dim` optional parameter is passed. What's the mechanism/way to understand more about this dispatch method (or is this just the PyTorch dispatcher)?
> is there a description of the new versus old style of batching rules?

Unfortunately no, not yet. [This is the doc](https://github.com/facebookresearch/functorch/blob/main/functorch/csrc/BatchRulesViews.cpp#L17) for the new style, but the old style was me hacking everything to work.

> I'm having a bit of trouble understanding what the following does and whether it's still needed in the new-style:

That is not needed in the new-style if you use the VMAP_SUPPORT macro.

> Is it right to assume that the optional<uint_64> parameter passed after every Tensor is the number of batch dimensions (e.g. if it is set to 2 then the first 2 dimensions of the Tensor are batch dimensions?).

The `optional<int64_t>` passed after every Tensor is "the index of the batch dimension" (not the number of batch dimensions!) The new-style batching rules assume that there is only a single batch dimension (but there is some magic somewhere else that allows the batching rule to operate on multiple batch dimensions).

> If it is empty then there are no batch dimensions?

Yes 

> Also if I return an empty batch dimension - then it appears to assume that there is 1 batch dimension (or this might be the original input number)?

You should be able to return `nullopt`

> Also there's implementations for squeeze and squeeze.dim - through trial-and-error I figured out that squeeze.dim is invoked when the dim optional parameter is passed. What's the mechanism/way to understand more about this dispatch method (or is this just the PyTorch dispatcher)?

This is just the PyTorch dispatcher, but the summary is:
- Look in [native_functions.yaml](https://github.com/pytorch/pytorch/blob/8fe32c9c1303446997f27e40e1a7c5645b1db367/aten/src/ATen/native/native_functions.yaml#L3899-L3911)
- Depending on how you invoke `squeeze` in python, one of the {squeeze, squeeze.dim} operators gets called. If one just calls `x.squeeze()`, it'll invoke the `squeeze` operator, if one calls `x.squeeze(1)`, it'll invoke the `squeeze.dim` operator.
- There's some logic in PyTorch's python bindings that *parses* the inputs to e.g. `squeeze` and selects which of {squeeze, squeeze.dim} actually gets called. If you're interested I can point you to that code (it is autogenerated so I can't link you it directly on github)





Fix patch:
diff --git a/functorch/csrc/BatchRulesViews.cpp b/functorch/csrc/BatchRulesViews.cpp
index a26b00f4e..0366e3cdd 100644
--- a/functorch/csrc/BatchRulesViews.cpp
+++ b/functorch/csrc/BatchRulesViews.cpp
@@ -189,6 +189,71 @@ const Tensor& resize__plumbing(
   return self;
 }
 
+std::tuple<Tensor, optional<int64_t>> squeeze_batching_rule(const Tensor& self, optional<int64_t> bdim) {
+  TORCH_INTERNAL_ASSERT(bdim.has_value());
+  // Special case for scalar arrays to replicate PyTorch behavior.
+  if (self.dim() == 1) {
+    return std::make_tuple(self.alias(), bdim);
+  }
+
+  // Manually calculate the output shape by eliding all dimensions of
+  // size 1 keeping track of where the batch index started and where it
+  // ended up moving to. We also ensure we do not drop the batch index.
+  auto shape = self.sizes();
+  DimVector squeezed_sizes;
+  bool before_batch_idx = true;
+  int64_t new_batch_idx = 0;
+  int64_t original_idx = 0;
+
+  for (auto it = shape.begin(); it != shape.end(); ++it) {
+    // Keep only dimensions != 1 and the batch dimension (irrespective of size).
+    if (*it != 1 || original_idx == bdim) {
+      squeezed_sizes.push_back(*it);
+      if (original_idx == bdim) {
+        before_batch_idx = false;
+      }
+      // Only increment for the dimensions that will be kept in the output.
+      if (before_batch_idx) {
+        ++new_batch_idx;
+      }
+    }
+    ++original_idx;
+  }
+
+  auto result = self.view(squeezed_sizes);
+  return std::make_tuple(result, c10::optional<int64_t>(new_batch_idx));
+}
+
+std::tuple<Tensor, optional<int64_t>> squeeze_dim_batching_rule(const Tensor& self, optional<int64_t> bdim, int64_t dim) {
+  TORCH_INTERNAL_ASSERT(bdim.has_value());
+  // Special case for scalar arrays to replicate PyTorch behavior.
+  if (self.dim() == 1) {
+    TORCH_CHECK(dim == 0, "Dimension is out of range (expected to be in range of [-1, 0], but got ", dim);
+    return std::make_tuple(self.alias(), bdim);
+  }
+
+  // Calculate the proper offset if dim is negative.
+  auto actual_dim = dim;
+  if (dim < 0) {
+    actual_dim = self.dim() + dim - 1;
+  }
+  if (actual_dim < bdim) {
+    // Since dimension to be squeezed is before the batch dimension pass as-is.
+    auto original_size = self.dim();
+    auto result = self.squeeze(actual_dim);
+    auto updated_batch_idx = *bdim;
+    if (result.dim() != original_size) {
+      // A column before batch dimension has been dropped so adjust accordingly.
+      --updated_batch_idx;
+    }
+    return std::make_tuple(result, optional<int64_t>(updated_batch_idx));
+  } else {
+    // Since dimension to be squeezed is after the batch dimension adjust by one to account
+    // for the original batch dimension. In this case batch dimension won't move.
+    return std::make_tuple(self.squeeze(actual_dim + 1), bdim);
+  }
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("diag", diag_batch_rule);
 
@@ -204,6 +269,8 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("_unsafe_view", _unsafe_view_batch_rule);
   VMAP_SUPPORT("unsqueeze", unsqueeze_batch_rule);
   m.impl("resize_", resize__plumbing);
+  VMAP_SUPPORT("squeeze", squeeze_batching_rule);
+  VMAP_SUPPORT("squeeze.dim", squeeze_dim_batching_rule);
 }
 
 }}
diff --git a/functorch/csrc/BatchingRegistrations.cpp b/functorch/csrc/BatchingRegistrations.cpp
index c9c637952..dac7101e4 100644
--- a/functorch/csrc/BatchingRegistrations.cpp
+++ b/functorch/csrc/BatchingRegistrations.cpp
@@ -448,42 +448,6 @@ Tensor& zero_inplace_batching_rule(Tensor &self) {
   return self;
 }
 
-Tensor squeeze_batching_rule(const Tensor& self) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.squeeze();
-  }
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto physical_sizes = self_physical.tensor().sizes();
-
-  // Don't squeeze the batch dims!
-  VmapDimVector squeezed_sizes;
-  int64_t num_batch_dims = self_physical.numBatchDims();
-  squeezed_sizes.insert(
-      squeezed_sizes.end(),
-      physical_sizes.begin(),
-      physical_sizes.begin() + num_batch_dims);
-  for (auto it = physical_sizes.begin() + num_batch_dims; it != physical_sizes.end(); ++it) {
-    if (*it != 1) {
-      squeezed_sizes.push_back(*it);
-    }
-  }
-
-  auto result = self_physical.tensor().view(squeezed_sizes);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
-Tensor squeeze_dim_batching_rule(const Tensor& self, int64_t dim) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.squeeze(dim);
-  }
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto dim_physical = self_physical.getPhysicalDim(dim);
-  auto result = self_physical.tensor().squeeze(dim_physical);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
 Tensor trace_batching_rule(const Tensor& self) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
@@ -1372,8 +1336,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   m.impl("slice.Tensor", slice_batching_rule);
   m.impl("split.Tensor", split_batching_rule);
   m.impl("split_with_sizes", split_with_sizes_batching_rule);
-  m.impl("squeeze", squeeze_batching_rule);
-  m.impl("squeeze.dim", squeeze_dim_batching_rule);
   m.impl("squeeze_.dim", squeeze_dim__batching_rule);
   m.impl("t", native::t); // composite wrt autograd
   // m.impl("trace", trace_batching_rule);


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index aa256912d..cd58d0a95 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -2106,16 +2106,36 @@ def test_slice(self):
              (torch.rand(3, 5, B0, B1, B2),), in_dims=2)
 
     def test_squeeze(self):
-        test = self._vmap_view_test
-        op = torch.squeeze
-        B0, B1 = 1, 11
-        test(op, (torch.rand(B0),))
-        test(op, (torch.rand(B0, 3, 5),))
-        test(op, (torch.rand(1, B0, 5),), in_dims=1)
-        test(op, (torch.rand(B0, 0, 1, 5, 1),))
-        test(op, (torch.rand(B0, 1, 1, 1, 1),))
-        test(vmap(op), (torch.rand(B0, B1, 1),))
-        test(vmap(op), (torch.rand(B1, 1, B0),), in_dims=2)
+        def verify_behavior(op, min_ndim=1):
+            test = self._vmap_view_test
+            B0, B1 = 1, 11
+            # These tests cannot be used with an operator that requires more
+            # than 1 dimension after batching.
+            if min_ndim <= 1:
+                test(op, (torch.rand(B0),))
+                test(op, (torch.rand(B1),))
+                test(vmap(op), (torch.rand(B0, B1, 1),))
+                test(vmap(op), (torch.rand(B1, 1, B0),), in_dims=2)
+            test(op, (torch.rand(B0, 3, 5),))
+            test(op, (torch.rand(1, B0, 5),), in_dims=1)
+            test(op, (torch.rand(B0, 0, 1, 5, 1),))
+            test(op, (torch.rand(B0, 1, 1, 1, 1),))
+            test(vmap(op), (torch.rand(B0, B1, 1, 3, 4),))
+            test(vmap(op), (torch.rand(B1, 1, B0, 4, 5),), in_dims=2)
+
+        verify_behavior(torch.squeeze)
+        verify_behavior(lambda x: torch.squeeze(x, dim=0), min_ndim=1)
+        verify_behavior(lambda x: torch.squeeze(x, dim=1), min_ndim=2)
+        verify_behavior(lambda x: torch.squeeze(x, dim=-1), min_ndim=2)
+        verify_behavior(lambda x: torch.squeeze(x, dim=-2), min_ndim=3)
+
+        msg = ""
+        try:
+            torch.squeeze(torch.rand(10), dim=1)
+        except IndexError as err:
+            msg = str(err)
+        with self.assertRaises(RuntimeError, msg=msg):
+            vmap(lambda x: torch.squeeze(x, dim=1))(torch.rand(10))
 
     def _test_mean_sum_dim(self, op):
         test = self._vmap_test
@@ -2905,7 +2925,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         # Each one of these is a bug
         vmap_fail = {
             '__getitem__',
-            'squeeze',
             'unfold',
             'nn.functional.linear',
         }