[port] expand to new api (#161)

* [port] expand to new api

* update code

* update code

* fix incorrect merge

* retrigger CI

Author: kshitij12345

functorch/csrc/BatchRulesViews.cpp

@@ -354,6 +354,40 @@ std::tuple<Tensor,optional<int64_t>> slice_backward_batch_rule(
   return std::make_tuple(std::move(result), 0);
 }
 
+std::tuple<Tensor, optional<int64_t>> expand_batch_rule(
+    const Tensor &self, optional<int64_t> self_bdim, IntArrayRef size, bool implicit)
+{
+  auto self_dim = self.dim();
+  TORCH_CHECK(static_cast<uint64_t>(self_dim - 1) <= size.size(),
+              "expand: the number of sizes provided (", size.size(), ") ",
+              "must be greater or equal to the number of dimensions in the tensor (", static_cast<uint64_t>(self_dim - 1), ")");
+
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto self_sizes = self_.sizes();
+  auto batch_size = self_sizes[0];
+
+  c10::SmallBuffer<int64_t, 5> size_(size.size() + 1);
+  size_[0] = batch_size;
+  std::copy(size.cbegin(), size.cend(), size_.begin() + 1);
+
+  // Here, we know we are expanding a (logical) tensor to a larger number
+  // of dimensions. We have to be careful because we can't call expand directly
+  // due to the presence of batch dimensions.
+  //
+  // As an example, let B0 be a batch dimension and consider expand(Tensor[B0, 3], [2, 3]).
+  // The result should be a tensor of size [B0, 2, 3].
+  // A physical view of size [B0, 3] can't directly be expanded to size [B0, 2, 3]
+  // so the strategy here is to view it first as a tensor of size [B0, 1, 3] and
+  // then expand.
+  auto extra_dims = size.size() - (self_dim - 1);
+  VmapDimVector view_shape(size_.size(), /*init_value*/1);
+  view_shape[0] = batch_size;
+  std::copy(self_sizes.cbegin() + 1, self_sizes.cend(),
+            view_shape.begin() + 1 + extra_dims);
+
+  return std::make_tuple(self_.view(view_shape).expand(size_, implicit), 0);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("diag", diag_batch_rule);
   VMAP_SUPPORT("chunk", chunk_batching_rule);
@@ -375,6 +409,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("diagonal_backward", diagonal_backward_batch_rule);
   VMAP_SUPPORT("select_backward", select_backward_batch_rule);
   VMAP_SUPPORT("slice_backward", slice_backward_batch_rule);
+  VMAP_SUPPORT("expand", expand_batch_rule);
 }
 
 }}

functorch/csrc/BatchingRegistrations.cpp

@@ -132,49 +132,6 @@ bool isPhysicalScalarTensor(const Tensor& logical_tensor) {
   return true;
 }
 
-Tensor expand_batching_rule(const Tensor& self, IntArrayRef size, bool implicit) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.expand(size, implicit);
-  }
-
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto size_physical = self_physical.getPhysicalShape(size);
-  auto self_physical_dim = self_physical.tensor().dim();
-
-  TORCH_CHECK((uint64_t)self_physical_dim <= size_physical.size(),
-       "expand: the number of sizes provided (", /*logical*/size.size(), ") ",
-       "must be greater or equal to the number of dimensions in the tensor (",
-       /*logical dim*/self.dim(), ")");
-
-  if ((uint64_t)self_physical_dim == size_physical.size()) {
-    auto result = self_physical.tensor().expand(size_physical, implicit);
-    return self_physical.getPhysicalToLogicalMap().apply(result);
-  }
-
-  TORCH_INTERNAL_ASSERT((uint64_t)self_physical_dim < size_physical.size());
-  // Here, we know we are expanding a (logical) tensor to a larger number
-  // of dimensions. We have to be careful because we can't call expand directly
-  // due to the presence of batch dimensions.
-  //
-  // As an example, let B0 be a batch dimension and consider expand(Tensor[B0, 3], [2, 3]).
-  // The result should be a tensor of size [B0, 2, 3].
-  // A physical view of size [B0, 3] can't directly be expanded to size [B0, 2, 3]
-  // so the strategy here is to view it first as a tensor of size [B0, 1, 3] and
-  // then expand.
-  auto self_physical_size = self_physical.tensor().sizes();
-  auto extra_dims = size_physical.size() - self_physical_dim;
-  VmapDimVector view_shape(size_physical.size(), 1);
-  std::copy(self_physical_size.begin(),
-            self_physical_size.begin() + self_physical.numBatchDims(),
-            view_shape.begin());
-  std::copy(self_physical_size.begin() + self_physical.numBatchDims(),
-            self_physical_size.end(),
-            view_shape.begin() + self_physical.numBatchDims() + extra_dims);
-  auto result = self_physical.tensor().view(view_shape).expand(size_physical, implicit);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
 std::vector<Tensor> chunk_batching_rule(const Tensor& self, int64_t chunks, int64_t dim) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
@@ -1001,7 +958,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   // m.impl("chunk", chunk_batching_rule);
   m.impl("tensor_split.sections", tensor_split_sections_batching_rule);
   m.impl("tensor_split.indices", tensor_split_indices_batching_rule);
-  m.impl("expand", expand_batching_rule);
   m.impl("movedim.intlist", movedim_batching_rule);
   m.impl("movedim.int", static_cast<Tensor(*)(const Tensor&,int64_t,int64_t)>(native::movedim)); // composite wrt autograd
   // NB: static_cast because there's another variant of narrow. However, we don't