Add swapdims op (#7659)

* add swapdims op

* add swapdims op

* revise

* revise

* revise

* revise op: swapaxes, swapdims, Transpose2dim

* revise op: swapaxes, swapdims, Transpose2dim

* revise Swapaxes and Swapdims

* fix conflic

* auto format by CI

* fix python/oneflow/framework/docstr/swapdims.py

* fix python/oneflow/framework/docstr/swapdims.py

* refine

Co-authored-by: chengyuma <chengyuma@github.com>
Co-authored-by: oneflow-ci-bot <ci-bot@oneflow.org>
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
Co-authored-by: Luyang <flowingsun007@163.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

docs/source/oneflow.rst

@@ -143,6 +143,7 @@ oneflow
             sqrt, 
             square,  
             swapaxes, 
+            swapdims, 
             tan, 
             tanh, 
             tensor, 

docs/source/tensor.rst

@@ -153,6 +153,7 @@ OneFlow Tensor Class
             stride, 
             sum,
             swapaxes, 
+            swapdims, 
             sub, 
             sub_, 
             tan, 

oneflow/core/functional/functional_api.yaml

@@ -330,10 +330,11 @@
   bind_python: True
 
 - name: "swapaxes"
-  signature:
-    [
-      "Tensor (Tensor input, Int32 dim0, Int32 dim1) => Swapaxes",
-    ]
+  signature: "Tensor (Tensor input, Int32 dim0, Int32 dim1) => Swapaxes"
+  bind_python: True
+
+- name: "swapdims"
+  signature: "Tensor (Tensor input, Int32 dim0, Int32 dim1) => Swapdims"
   bind_python: True
 
 - name: "permute"

oneflow/core/functional/impl/math_functor.cpp

@@ -732,9 +732,11 @@ class Transpose2dimFunctor {
     if (dim1 < 0) { dim_1 += ndim; }
 
     CHECK_OR_RETURN(dim_0 >= 0 && dim0 < ndim)
-        << "Invalid dim0:" << dim_0 << " len(shape):" << ndim;
+        << "Dimension out of range (expected to be in range of [" << -ndim << ", " << ndim - 1
+        << "], but got " << dim_0 << ")";
     CHECK_OR_RETURN(dim_1 >= 0 && dim1 < ndim)
-        << "Invalid dim1:" << dim_1 << " len(shape):" << ndim;
+        << "Dimension out of range (expected to be in range of [" << -ndim << ", " << ndim - 1
+        << "], but got " << dim_1 << ")";
     for (int32_t i = 0; i < ndim; ++i) { permute.emplace_back(i); }
     std::swap(permute[dim_0], permute[dim_1]);
 
@@ -793,28 +795,6 @@ class AsStridedGradFunctor {
   std::shared_ptr<OpExpr> op_;
 };
 
-class SwapaxesFunctor {
- public:
-  SwapaxesFunctor() {}
-  Maybe<Tensor> operator()(const std::shared_ptr<one::Tensor>& x, const int32_t dim0,
-                           const int32_t dim1) const {
-    const int64_t ndim = x->shape()->NumAxes();
-    int32_t dim_0 = dim0;
-    int32_t dim_1 = dim1;
-
-    if (dim0 < 0) { dim_0 += ndim; }
-    if (dim1 < 0) { dim_1 += ndim; }
-
-    CHECK_OR_RETURN(dim_0 >= 0 && dim0 < ndim)
-        << "Dimension out of range (expected to be in range of [" << -ndim << ", " << ndim - 1
-        << "], but got " << dim_0 << ")";
-    CHECK_OR_RETURN(dim_1 >= 0 && dim1 < ndim)
-        << "Dimension out of range (expected to be in range of [" << -ndim << ", " << ndim - 1
-        << "], but got " << dim_1 << ")";
-    return Transpose2dim(x, dim0, dim1);
-  }
-};
-
 class ArangeFunctor {
  public:
   ArangeFunctor() { op_ = CHECK_JUST(one::OpBuilder("arange").Output("out").Build()); }
@@ -2804,7 +2784,8 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
   m.add_functor<TransposeFunctor>("Permute");
   m.add_functor<AsStridedFunctor>("AsStrided");
   m.add_functor<AsStridedGradFunctor>("AsStridedGrad");
-  m.add_functor<SwapaxesFunctor>("Swapaxes");
+  m.add_functor<Transpose2dimFunctor>("Swapaxes");
+  m.add_functor<Transpose2dimFunctor>("Swapdims");
   m.add_functor<ArangeFunctor, Arange2Functor>("Arange");
   m.add_functor<ConsistentArangeFunctor, ConsistentArange2Functor>("ConsistentArange");
   m.add_functor<CastFunctor>("Cast");

python/oneflow/__init__.py

@@ -181,6 +181,7 @@ def is_deprecated(func_or_class):
 from oneflow._C import cumsum
 from oneflow._C import cumprod
 from oneflow._C import swapaxes
+from oneflow._C import swapdims
 from oneflow._C import t
 from oneflow._C import masked_fill
 from oneflow._C import equal

python/oneflow/framework/docstr/__init__.py

@@ -61,6 +61,7 @@
 from .index_select import *
 from .sort import *
 from .is_floating_point import *
+from .swapdims import *
 from .where import *
 from .einsum import *
 from .oneflow import *

python/oneflow/framework/docstr/swapaxes.py

@@ -17,8 +17,10 @@
 from oneflow.framework.docstr.utils import add_docstr
 
 add_docstr(
-    oneflow.swapaxes,
-    """This function is equivalent to NumPy’s swapaxes function.
+    oneflow._C.swapaxes,
+    """swapaxes(input, axis0, axis1) -> Tensor
+    
+    This function is equivalent to NumPy’s swapaxes function.
 
     For example:
 

python/oneflow/framework/docstr/swapdims.py

@@ -0,0 +1,53 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import oneflow
+from oneflow.framework.docstr.utils import add_docstr
+
+add_docstr(
+    oneflow._C.swapdims,
+    """
+    swapdims(input, dim0, dim1) -> Tensor
+
+    This function is equivalent to torch’s swapdims function.
+
+    For example:
+
+    .. code-block:: python
+    
+        >>> import oneflow as flow
+
+        >>> x = flow.tensor([[[0,1],[2,3]],[[4,5],[6,7]]])
+        >>> x
+        tensor([[[0, 1],
+                 [2, 3]],
+        <BLANKLINE>
+                [[4, 5],
+                 [6, 7]]], dtype=oneflow.int64)
+        >>> flow.swapdims(x, 0, 1)
+        tensor([[[0, 1],
+                 [4, 5]],
+        <BLANKLINE>
+                [[2, 3],
+                 [6, 7]]], dtype=oneflow.int64)
+        >>> flow.swapdims(x, 0, 2)
+        tensor([[[0, 4],
+                 [2, 6]],
+        <BLANKLINE>
+                [[1, 5],
+                 [3, 7]]], dtype=oneflow.int64)
+
+    """,
+)

python/oneflow/framework/docstr/tensor.py

@@ -756,6 +756,13 @@
     """,
 )
 
+add_docstr(
+    oneflow.Tensor.swapdims,
+    """
+    See :func:`oneflow.swapdims`
+    """,
+)
+
 add_docstr(
     oneflow.Tensor.cast,
     """

python/oneflow/framework/tensor.py

@@ -389,6 +389,10 @@ def _swapaxes(self, dim0, dim1):
     return flow._C.swapaxes(self, dim0, dim1)
 
 
+def _swapdims(self, dim0, dim1):
+    return flow._C.swapdims(self, dim0, dim1)
+
+
 def _cast(self, dtype):
     return flow.cast(self, dtype)
 
@@ -1213,6 +1217,7 @@ def RegisterMethods():
     Tensor.unbind = _unbind
     Tensor.squeeze = _squeeze
     Tensor.swapaxes = _swapaxes
+    Tensor.swapdims = _swapdims
     Tensor.unfold = _unfold
     Tensor.narrow = _narrow
     Tensor.unsqueeze = _unsqueeze

python/oneflow/test/modules/test_swapdims.py

@@ -0,0 +1,47 @@
+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import unittest
+from collections import OrderedDict
+
+import numpy as np
+
+from oneflow.test_utils.automated_test_util import *
+from oneflow.test_utils.test_util import GenArgList
+
+import oneflow as flow
+import oneflow.unittest
+
+
+@flow.unittest.skip_unless_1n1d()
+class Testswapdims(flow.unittest.TestCase):
+    @autotest(check_graph=True)
+    def test_swapdims_flow_with_random_data(test_case):
+        device = random_device()
+        x = random_tensor(ndim=3).to(device)
+        y = torch.swapdims(x, np.random.randint(0, 3), np.random.randint(0, 3))
+        return y
+
+    @autotest(check_graph=True)
+    def test_swapdims_flow_with_random_data2(test_case):
+        device = random_device()
+        x = random_tensor(ndim=4).to(device)
+        y = torch.swapdims(x, np.random.randint(0, 4), np.random.randint(0, 4))
+        return y
+
+
+if __name__ == "__main__":
+    unittest.main()

python/oneflow/test/tensor/test_tensor_part_2.py

@@ -915,6 +915,14 @@ def test_tensor_swapaxes(test_case):
         y = x.swapaxes(random(0, 2).to(int), random(0, 2).to(int))
         return y
 
+    @flow.unittest.skip_unless_1n1d()
+    @autotest(n=5, check_graph=True)
+    def test_tensor_swapdimst(test_case):
+        device = random_device()
+        x = random_tensor(ndim=3).to(device)
+        y = x.swapdims(random(0, 3).to(int), random(0, 3).to(int))
+        return y
+
 
 if __name__ == "__main__":
     unittest.main()