[port] view to new api (#164)

* [port] view to new api

* handle no batch dim case

* forward the correct bdim

* refactor non bdim code

* address review

Author: kshitij12345

functorch/csrc/BatchRulesViews.cpp

@@ -372,6 +372,18 @@ std::tuple<Tensor,optional<int64_t>> slice_backward_batch_rule(
   return std::make_tuple(std::move(result), 0);
 }
 
+std::tuple<Tensor, optional<int64_t>> view_batching_rule(
+    const Tensor &self, optional<int64_t> self_bdim, IntArrayRef size)
+{
+  TORCH_INTERNAL_ASSERT(self_bdim.has_value());
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  VmapDimVector size_(size.size() + 1);
+  // copy batch size
+  size_[0] = self_.size(0);
+  std::copy(size.cbegin(), size.cend(), size_.begin() + 1);
+  return std::make_tuple(self_.view(size_), 0);
+}
+
 std::tuple<Tensor, optional<int64_t>> expand_batch_rule(
     const Tensor &self, optional<int64_t> self_bdim, IntArrayRef size, bool implicit)
 {
@@ -428,6 +440,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("view", view_batching_rule);
   VMAP_SUPPORT("expand", expand_batch_rule);
 }
 

functorch/csrc/BatchingRegistrations.cpp

@@ -407,17 +407,6 @@ Tensor contiguous_batching_rule(const Tensor& self, MemoryFormat memory_format)
   return physical_view.getPhysicalToLogicalMap().apply(result);
 }
 
-Tensor view_batching_rule(const Tensor& self, IntArrayRef size) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.view(size);
-  }
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto size_physical = self_physical.getPhysicalShape(size);
-  auto result = self_physical.tensor().view(size_physical);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
 Tensor view_as_complex_batching_rule(const Tensor& self) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
@@ -952,7 +941,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   m.impl("unbind.int", unbind_batching_rule);
   m.impl("unfold", unfold_batching_rule);
   m.impl("unsqueeze_", unsqueeze__batching_rule);
-  m.impl("view", view_batching_rule);
   m.impl("view_as", native::view_as); // composite wrt autograd
 
   m.impl("addmm", addmm_batching_rule);