(fix CI) Add batch rule for split.sizes (#952)

* add batch rule for split.sizes

* switch to doing split.sizes as a decomposition

Author: Samantha Andow

functorch/csrc/BatchRulesDecompositions.cpp

@@ -181,6 +181,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   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/csrc/LegacyBatchingRegistrations.cpp

@@ -270,11 +270,11 @@ std::vector<Tensor> split_batching_rule(const Tensor& self, int64_t split_size,
 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;
 }