convert movedim to new api (#211)

* convery movedim to new api

* move to view batching rules

* remove automatically unbatched

* convery movedim to new api

* move to view batching rules

* remove automatically unbatched

* remove unused

* remove merge typo

Author: Padarn

functorch/csrc/BatchRulesHelper.cpp

@@ -52,6 +52,22 @@ int64_t getPhysicalDim(const Tensor& tensor, bool has_batch_dim, int64_t logical
   return wrapped_dim;
 }
 
+VmapDimVector getPhysicalDims(const Tensor& tensor, bool has_batch_dim, IntArrayRef logical_dims) {
+  // NB: assumes the batch dim is at the front of the tensor
+  optional<int64_t> bdim = has_batch_dim ? optional<int64_t>(0) : nullopt;
+  auto rank = rankWithoutBatchDim(tensor, bdim);
+  VmapDimVector result;
+  result.reserve(logical_dims.size());
+  for (auto d : logical_dims){
+    if (has_batch_dim) {
+      result.push_back(maybe_wrap_dim(d, rank)+1);
+    } else {
+      result.push_back(maybe_wrap_dim(d, rank));
+    } 
+  }
+  return result;
+}
+
 Tensor maybePadToLogicalRank(const Tensor& tensor, optional<int64_t> has_bdim, int64_t logical_rank) {
   if (!has_bdim) {
     return tensor;

functorch/csrc/BatchRulesHelper.h

@@ -27,6 +27,7 @@ int64_t rankWithoutBatchDim(const Tensor& tensor, optional<int64_t> maybe_batch_
 int64_t numelWithoutBatchDim(const Tensor& tensor, optional<int64_t> maybe_batch_dim);
 optional<int64_t> valIfNonempty(optional<int64_t> maybe_empty, int64_t new_val);
 int64_t getPhysicalDim(const Tensor& tensor, bool has_batch_dim, int64_t logical_dim);
+VmapDimVector getPhysicalDims(const Tensor& tensor, bool has_batch_dim, IntArrayRef logical_dims);
 
 void vmapIncompatibleInplaceError(const char* schema_name);
 

functorch/csrc/BatchRulesViews.cpp

@@ -446,6 +446,13 @@ std::tuple<Tensor, optional<int64_t>> unfold_batch_rule(
   return std::make_tuple(result, 0);
 }
 
+std::tuple<Tensor, optional<int64_t>> movedim_batch_rule(const Tensor& self, optional<int64_t> self_bdim, IntArrayRef source, IntArrayRef destination) {
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto source_ = getPhysicalDims(self_, self_bdim.has_value(), source);
+  auto destination_ = getPhysicalDims(self_, self_bdim.has_value(), destination);
+  return std::make_tuple(self_.movedim(source_, destination_), 0);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("diag", diag_batch_rule);
   VMAP_SUPPORT("chunk", chunk_batching_rule);
@@ -471,6 +478,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("view", view_batching_rule);
   VMAP_SUPPORT("expand", expand_batch_rule);
   VMAP_SUPPORT("unfold", unfold_batch_rule);
-}
+  VMAP_SUPPORT("movedim.intlist", movedim_batch_rule);
+} 
 
 }}

functorch/csrc/BatchingRegistrations.cpp

@@ -308,18 +308,6 @@ Tensor slice_backward_batching_rule(const Tensor& grad, IntArrayRef input_sizes,
   return grad_physical.getPhysicalToLogicalMap().apply(grad_input);
 }
 
-Tensor movedim_batching_rule(const Tensor& self, IntArrayRef source, IntArrayRef destination) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return at::movedim(self, source, destination);
-  }
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto source_physical = self_physical.getPhysicalDims(source);
-  auto destination_physical = self_physical.getPhysicalDims(destination);
-  auto result = at::movedim(self_physical.tensor(), source_physical, destination_physical);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
 std::vector<Tensor> split_batching_rule(const Tensor& self, int64_t split_size, int64_t dim) {
   if (!participatesInCurrentLevel(self)) {
     c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
@@ -884,8 +872,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("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
   // want to support the other variant yet bc it isn't documented...
   m.impl("numpy_T", native::numpy_T); // composite wrt autograd