Add torch.movedim

aten/src/ATen/native/TensorShape.cpp

@@ -1693,4 +1693,46 @@ Tensor& diag_cpu_out(Tensor &result, const Tensor& self, int64_t dimension) {
   return result;
 }
 
+Tensor moveaxis(const Tensor& self, IntArrayRef src, IntArrayRef dst) {
+  TORCH_CHECK(src.size() == dst.size(), "moveaxis: Invalid source or destination dims: src (",
+              src.size(), " dims ) should contain the same number of dims as dst (", dst.size(), "dims)");
+  auto normalized_src = src.vec();
+  maybe_wrap_dims(normalized_src, self.dim());
+  auto normalized_dst = dst.vec();
+  maybe_wrap_dims(normalized_dst, self.dim());
+
+  auto it_src = std::unique(normalized_src.begin(), normalized_src.end());
+  TORCH_CHECK(it_src == normalized_src.end(), "moveaxis: repeated axis in `src` (", src, ")");
+  auto it_dst = std::unique(normalized_dst.begin(), normalized_dst.end());
+  TORCH_CHECK(it_dst == normalized_dst.end(), "moveaxis: repeated axis in `dst` (", dst, ")");
+
+  std::vector<int64_t> order, source_dims, destination_dims;
+  order.resize(self.dim());
+  source_dims.resize(self.dim());
+  destination_dims.resize(self.dim());
+
+  std::iota(source_dims.begin(), source_dims.end(), 0);
+  std::iota(destination_dims.begin(), destination_dims.end(), 0);
+
+  for (int64_t i = 0; i < src.size(); ++i) {
+      order[normalized_dst[i]] = normalized_src[i];
+      source_dims[normalized_src[i]] = -1;
+      destination_dims[normalized_dst[i]] = -1;
+  }
+
+  auto source_iter = std::remove(source_dims.begin(), source_dims.end(), -1);
+  auto destination_iter = std::remove(destination_dims.begin(), destination_dims.end(), -1);
+
+  int64_t rest_dim = self.dim() - src.size();
+  for (int64_t i = 0; i < rest_dim; ++i) {
+      order[destination_dims[i]] = source_dims[i];
+  }
+
+  return self.permute(order);
+}
+
+Tensor moveaxis(const Tensor& self, int64_t src, int64_t dst) {
+  return at::moveaxis(self, IntArrayRef{src}, IntArrayRef{dst});
+}
+
 }} // at::native

aten/src/ATen/native/native_functions.yaml

@@ -2159,6 +2159,14 @@
   use_c10_dispatcher: full
   variants: method  # This is method-only to match the previous tensor API. In the future we could make this a function too.
 
+- func: moveaxis.intlist(Tensor self, int[] src, int[] dst) -> Tensor
+  use_c10_dispatcher: full
+  variants: function
+
+- func: moveaxis.int(Tensor self, int src, int dst) -> Tensor
+  use_c10_dispatcher: full
+  variants: function
+
 # Only exposed from C++ -- in Python,
 # we expose it as an attribute `T`, not a function.
 #

test/test_autograd.py

@@ -6510,6 +6510,16 @@ def test_strided_leaf_grad_layout(self, device):
         (c * d).sum().backward()
         self.assertEqual(c.grad.stride(), (2, 1))
 
+    def test_moveaxis(self, device):
+        x = torch.randn(4, 3, 2, 1, dtype=torch.double, requires_grad=True)
+
+        # Positive axis
+        gradcheck(lambda x: torch.moveaxis(x, (0, 1, 2, 3), (3, 2, 1, 0)), x)
+        gradgradcheck(lambda x: torch.moveaxis(x, (0, 1, 2, 3), (3, 2, 1, 0)), x)
+
+        # Negative axis
+        gradcheck(lambda x: torch.moveaxis(x, (0, -1, -2, -3), (-3, -2, -1, -0)), x)
+        gradgradcheck(lambda x: torch.moveaxis(x, (0, -1, -2, -3), (-3, -2, -1, -0)), x)
 
 class TestMultithreadAutograd(TestCase):
     def _run_py_multithread_fn(self, fn, args=(), num_threads=10, kwargs=None):

test/test_torch.py

@@ -18121,6 +18121,57 @@ def test_large_linspace(self, device, dtype):
         x = torch.linspace(start, end, steps, dtype=dtype, device=device)
         self.assertGreater(x[1] - x[0], (end - start) / steps)
 
+    @dtypes(torch.int64, torch.float, torch.complex128)
+    def test_moveaxis_invalid(self, device, dtype):
+        shape = self._rand_shape(4, min_size=5, max_size=10)
+        x = self._generate_input(shape, dtype, device, False)
+
+        # Invalid `src` and `dst` dimension
+        with self.assertRaisesRegex(IndexError, "Dimension out of range"):
+            torch.moveaxis(x, 5, 0)
+
+        with self.assertRaisesRegex(IndexError, "Dimension out of range"):
+            torch.moveaxis(x, 0, 5)
+
+        # Mismatch in size of `src` and `dst`
+        with self.assertRaisesRegex(RuntimeError, "moveaxis: Invalid source or destination dims:"):
+            torch.moveaxis(x, (1, 0), (0, ))
+
+        with self.assertRaisesRegex(RuntimeError, "moveaxis: repeated axis in `src`"):
+            torch.moveaxis(x, (0, 0), (0, 1))
+
+        with self.assertRaisesRegex(RuntimeError, "moveaxis: repeated axis in `dst`"):
+            torch.moveaxis(x, (0, 1), (1, 1))
+
+    @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
+    @dtypes(torch.int64, torch.float, torch.complex128)
+    def test_moveaxis(self, device, dtype):
+        for nd in range(5):
+            shape = self._rand_shape(nd, min_size=5, max_size=10)
+            x = self._generate_input(shape, dtype, device, with_extremal=False)
+            for random_negative in [True, False]:
+                for src_dim, dst_dim in permutations(range(nd), r=2):
+                    if random_negative:
+                        src_dim = src_dim - nd
+                    # Integer Inputs
+                    torch_fn = partial(torch.moveaxis, src=src_dim, dst=dst_dim)
+                    np_fn = partial(np.moveaxis, source=src_dim, destination=dst_dim)
+                    self.compare_with_numpy(torch_fn, np_fn, x, device=None, dtype=None)
+
+                if nd > 0:
+                    for src_sequence in permutations(range(nd), r=random.randint(1, nd)):
+                        # Randomly change a dim to a negative dim representation of itself.
+                        if random_negative:
+                            random_idx = random.randint(0, len(src_sequence) - 1)
+                            src_sequence = list(src_sequence)
+                            src_sequence[random_idx] = src_sequence[random_idx] - nd
+                            src_sequence = tuple(src_sequence)
+                        # Sequence Inputs
+                        dst_sequence = tuple(random.sample(range(nd), len(src_sequence)))
+                        torch_fn = partial(torch.moveaxis, src=src_sequence, dst=dst_sequence)
+                        np_fn = partial(np.moveaxis, source=src_sequence, destination=dst_sequence)
+                        self.compare_with_numpy(torch_fn, np_fn, x, device=None, dtype=None)
+
 # NOTE [Linspace+Logspace precision override]
 # Our Linspace and logspace torch.half CUDA kernels are not very precise.
 # Since linspace/logspace are deterministic, we can compute an expected

torch/_overrides.py

@@ -407,6 +407,7 @@ def get_testing_overrides():
                            dropout, train, bidirectional, batch_sizes, dropout_state: -1),
         torch.mm: lambda input, mat2, out=None: -1,
         torch.mode: lambda input: -1,
+        torch.moveaxis: lambda input, src, dst: -1,
         torch.mul: lambda input, other, out=None: -1,
         torch.multinomial: lambda input, num_samples, replacement=False, out=None: -1,
         torch.mv: lambda input, vec, out=None: -1,

torch/_torch_docs.py

@@ -4365,6 +4365,44 @@ def merge_dicts(*dicts):
             [1.0311, 0.3901, 0.5049]])
 """)
 
+add_docstr(torch.moveaxis,
+           r"""
+moveaxis(input, src, dst) -> Tensor
+
+Move axes of an array to new positions.
+Other axes remain in their original order.
+
+Args:
+    input (Tensor)
+    src (int or tuple of ints): Original positions of the axes to move. These must be unique.
+    dst (int or tuple of ints): Destination positions for each of the original axes. These must also be unique.
+
+Example::
+
+    >>> a = torch.randn(3,2,1)
+    >>> a
+    tensor([[[-0.3362],
+            [-0.8437]],
+
+            [[-0.9627],
+            [ 0.1727]],
+
+            [[ 0.5173],
+            [-0.1398]]])
+    >>> torch.moveaxis(a, 1, 0)
+    tensor([[[-0.3362],
+            [-0.9627],
+            [ 0.5173]],
+
+            [[-0.8437],
+            [ 0.1727],
+            [-0.1398]]])
+    >>> torch.moveaxis(a, (1, 2), (0, 1))
+    tensor([[[-0.3362, -0.9627,  0.5173]],
+
+            [[-0.8437,  0.1727, -0.1398]]])
+""")
+
 add_docstr(torch.narrow,
            r"""
 narrow(input, dim, start, length) -> Tensor