add scatter.reduce batching rule (#188)

* add scatter.reduce batching rule

* minor update

Author: kshitij12345

functorch/csrc/BatchRulesScatterOps.cpp

@@ -186,11 +186,15 @@ int64_t bdim_size(
   TORCH_INTERNAL_ASSERT(false);
 }
 
-std::tuple<Tensor,optional<int64_t>> scatter_value_batch_rule(
+namespace {
+
+template<typename Func, typename ...Args>
+std::tuple<Tensor,optional<int64_t>> scatter_batch_rule(
+    Func f,
     const Tensor& self, optional<int64_t> self_bdim,
     int64_t dim,
     const Tensor& index, optional<int64_t> index_bdim,
-    const Scalar& value) {
+    const Scalar& value, Args... args) {
   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);
@@ -208,23 +212,21 @@ std::tuple<Tensor,optional<int64_t>> scatter_value_batch_rule(
   index_ = ensure_has_bdim(index_, index_bdim.has_value(), batch_size);
   auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
 
-  auto result = at::scatter(self_, physical_dim, index_, value);
-  // result should have same rank as self
+  auto result = f(self_, physical_dim, index_, value, args...);
+  // result should have same shape as self
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
   return std::make_tuple(result, 0);
 }
 
-namespace {
-
-template <typename Func>
+template <typename Func, typename ...Args>
 inline std::tuple<Tensor,optional<int64_t>> scatter_batch_rule(
     Func f,
     const Tensor& self, optional<int64_t> self_bdim,
     int64_t dim,
     const Tensor& index, optional<int64_t> index_bdim,
-    const Tensor& src, optional<int64_t> src_bdim) {
+    const Tensor& src, optional<int64_t> src_bdim, Args... args) {
   auto self_logical_rank = rankWithoutBatchDim(self, self_bdim);
   auto index_logical_rank = rankWithoutBatchDim(index, index_bdim);
   auto src_logical_rank = rankWithoutBatchDim(src, src_bdim);
@@ -248,8 +250,8 @@ inline std::tuple<Tensor,optional<int64_t>> scatter_batch_rule(
   src_ = ensure_has_bdim(src_, src_bdim.has_value(), batch_size);
   auto physical_dim = getPhysicalDim(self_, /*has_batch_dim*/true, dim);
 
-  auto result = f(self_, physical_dim, index_, src_);
-  // result should have same rank as self
+  auto result = f(self_, physical_dim, index_, src_, args...);
+  // result should have same shape as self
   if (self_logical_rank == 0) {
     result = result.squeeze(-1);
   }
@@ -258,6 +260,15 @@ inline std::tuple<Tensor,optional<int64_t>> scatter_batch_rule(
 
 } // namespace
 
+std::tuple<Tensor,optional<int64_t>> scatter_value_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    int64_t dim,
+    const Tensor& index, optional<int64_t> index_bdim,
+    const Scalar& value) {
+  return scatter_batch_rule(ATEN_FN2(scatter, value),
+                            self, self_bdim, dim, index, index_bdim, value);
+}
+
 std::tuple<Tensor,optional<int64_t>> scatter_src_batch_rule(
     const Tensor& self, optional<int64_t> self_bdim,
     int64_t dim,
@@ -276,6 +287,28 @@ std::tuple<Tensor,optional<int64_t>> scatter_add_batch_rule(
                             self, self_bdim, dim, index, index_bdim, src, src_bdim);
 }
 
+std::tuple<Tensor,optional<int64_t>> scatter_reduce_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    int64_t dim,
+    const Tensor& index, optional<int64_t> index_bdim,
+    const Tensor& src, optional<int64_t> src_bdim,
+    const c10::string_view reduce) {
+  using scatter_reduce_value_sig = Tensor (*)(const Tensor&, int64_t, const Tensor&, const Tensor&, const c10::string_view reduce);
+  return scatter_batch_rule(ATEN_FN2(scatter, reduce),
+                            self, self_bdim, dim, index, index_bdim, src, src_bdim, reduce);
+}
+
+std::tuple<Tensor,optional<int64_t>> scatter_value_reduce_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    int64_t dim,
+    const Tensor& index, optional<int64_t> index_bdim,
+    const Scalar& src,
+    const c10::string_view reduce) {
+  using scatter_reduce_value_sig = Tensor (*)(const Tensor&, int64_t, const Tensor&, const Scalar&, const c10::string_view reduce);
+  return scatter_batch_rule(ATEN_FN2(scatter, value_reduce),
+                            self, self_bdim, dim, index, index_bdim, src, reduce);
+}
+
 std::tuple<Tensor,optional<int64_t>> gather_batch_rule(
     const Tensor& self, optional<int64_t> self_bdim,
     int64_t dim,
@@ -400,6 +433,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("scatter.reduce", scatter_reduce_batch_rule);
+  VMAP_SUPPORT("scatter.value_reduce", scatter_value_reduce_batch_rule);
   VMAP_SUPPORT("index_select", index_select_batch_rule);
 
 }

test/test_vmap.py

@@ -3089,7 +3089,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('resize_as_'),
         xfail('resolve_conj'),
         xfail('resolve_neg'),
-        xfail('scatter'),
         xfail('take'),
         xfail('take_along_dim'),
         xfail('tensor_split'),