commit 65e8d2209c0c89adff922fa5390aff66c4c3ce0d
author Samantha Andow <samdow@fb.com> Mon Jul 11 13:06:21 2022 -0400
committer Jon Janzen <jonjanzen@fb.com> Thu Jul 21 13:41:37 2022 -0700
tree 36079a5197d3c0eaa25055a07e341d5ae7cefd9d
parent 8445222264f582c68c2b4caf6b3fb74b253b14b8

[functorch] (fix CI) Add batch rule for split.sizes (pytorch/functorch#952)

* add batch rule for split.sizes

* switch to doing split.sizes as a decomposition

functorch/functorch/csrc/BatchRulesDecompositions.cpp

diff --git a/functorch/functorch/csrc/BatchRulesDecompositions.cpp b/functorch/functorch/csrc/BatchRulesDecompositions.cpp
index b84940b..aff7acf 100644
--- a/functorch/functorch/csrc/BatchRulesDecompositions.cpp
+++ b/functorch/functorch/csrc/BatchRulesDecompositions.cpp
@@ -181,6 +181,7 @@
   OP_DECOMPOSE(special_multigammaln);
   OP_DECOMPOSE(special_polygamma);
   OP_DECOMPOSE(special_softmax);
+  OP_DECOMPOSE2(split, sizes);
   OP_DECOMPOSE(square);
   OP_DECOMPOSE(numpy_T);
   OP_DECOMPOSE(reshape_as);

functorch/functorch/csrc/LegacyBatchingRegistrations.cpp

diff --git a/functorch/functorch/csrc/LegacyBatchingRegistrations.cpp b/functorch/functorch/csrc/LegacyBatchingRegistrations.cpp
index 0bdbca5..8181174 100644
--- a/functorch/functorch/csrc/LegacyBatchingRegistrations.cpp
+++ b/functorch/functorch/csrc/LegacyBatchingRegistrations.cpp
@@ -270,11 +270,11 @@
 std::vector<Tensor> split_with_sizes_batching_rule(const Tensor& self, IntArrayRef split_sizes, int64_t dim) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return at::split_with_sizes(self, split_sizes, dim);
+    return split_with_sizes(self, split_sizes, dim);
   }
   auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
   auto dim_physical = self_physical.getPhysicalDim(dim);
-  auto result = at::split_with_sizes(self_physical.tensor(), split_sizes, dim_physical);
+  auto result = split_with_sizes(self_physical.tensor(), split_sizes, dim_physical);
   self_physical.getPhysicalToLogicalMap().applyInplace(result);
   return result;
 }