Added index_select batching rule (#183)

* Added index_select batching rule
Description:
- Added index_select batching rule
- Updated tests

Note: index_select_backward can not done due to its composite nature and `index_add_` usage

* Attempt to fix failing tests

* Fixed issue with tests

* Fixed nit comment

Author: vfdev-5

functorch/csrc/BatchRulesScatterOps.cpp

@@ -219,7 +219,7 @@ std::tuple<Tensor,optional<int64_t>> scatter_value_batch_rule(
   auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
 
   auto result = at::scatter(self_, physical_dim, index_, value);
-  // result should have same shape as self
+  // result should have same rank as self
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
@@ -259,7 +259,7 @@ inline std::tuple<Tensor,optional<int64_t>> scatter_batch_rule(
   auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
 
   auto result = f(self_, physical_dim, index_, src_);
-  // result should have same shape as self
+  // result should have same rank as self
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
@@ -309,7 +309,7 @@ std::tuple<Tensor,optional<int64_t>> gather_batch_rule(
   auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
 
   auto result = at::gather(self_, physical_dim, index_, sparse_grad);
-  // result should have same shape as index
+  // result should have same rank as index
   if (index_logical_rank == 0) {
     result = result.squeeze(-1);
   }
@@ -346,7 +346,56 @@ std::tuple<Tensor,optional<int64_t>> gather_backward_batch_rule(
 
   auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
   auto result = at::gather_backward(grad_, self_, physical_dim, index_, sparse_grad);
-  // result should has same shape as self
+  // result should has same rank as self
+  if (self_logical_rank == 0) {
+    result = result.squeeze(-1);
+  }
+  return std::make_tuple(result, 0);
+}
+
+std::tuple<Tensor, optional<int64_t>> index_select_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    int64_t dim,
+    const Tensor& index, optional<int64_t> index_bdim) {
+
+  auto self_logical_rank = rankWithoutBatchDim(self, self_bdim);
+  auto index_logical_rank = rankWithoutBatchDim(index, index_bdim);
+  auto batch_size = bdim_size(self, self_bdim, index, index_bdim);
+
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto index_ = moveBatchDimToFront(index, index_bdim);
+
+  if (self_logical_rank == 0) {
+    self_ = self_.unsqueeze(-1);
+  }
+  if (index_logical_rank == 0) {
+    index_ = index_.unsqueeze(-1);
+  }
+  self_ = ensure_has_bdim(self_, self_bdim.has_value(), batch_size);
+  index_ = ensure_has_bdim(index_, index_bdim.has_value(), batch_size);
+  auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
+
+  if (index_.dim() < self_.dim()) {
+    // setup new_index_shape as [BS, 1, ..., le, ..., 1]
+    // to reshape index_
+    auto le = index_.size(1);  // get non-batch size of index tensor
+    {
+      VmapDimVector new_index_shape(self_.dim(), 1);
+      new_index_shape[0] = self_.size(0); // set up batch size
+      new_index_shape[physical_dim] = le;
+      index_ = index_.reshape(new_index_shape);
+    }
+    // Now apply expand to index_
+    {
+      auto self_shape = self_.sizes();
+      VmapDimVector new_index_shape = {self_shape.begin(), self_shape.end()};
+      new_index_shape[physical_dim] = le;
+      index_ = index_.expand(new_index_shape);
+    }
+  }
+
+  auto result = at::gather(self_, physical_dim, index_);
+  // result should have same rank as self
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
@@ -361,6 +410,8 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("scatter.value", scatter_value_batch_rule);
   VMAP_SUPPORT("scatter.src", scatter_src_batch_rule);
   VMAP_SUPPORT("scatter_add", scatter_add_batch_rule);
+  VMAP_SUPPORT("index_select", index_select_batch_rule);
+
 }
 
 }}

test/test_vmap.py

@@ -2994,7 +2994,7 @@ class TestVmapOperatorsOpInfo(TestCase):
         xfail('block_diag'),
         xfail('nn.functional.dropout'),
 
-        # entries in here need don't work and need to be fixed.
+        # entries in here don't work and need to be fixed.
         # Each one of these is a bug
         xfail('unfold'),
         xfail('svd', device_type='cuda'),
@@ -3041,7 +3041,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('index_copy'),
         xfail('index_fill'),
         xfail('index_put'),
-        xfail('index_select'),
         xfail('isin'),
         xfail('linalg.cholesky'),
         xfail('linalg.eigvals'),