Bug fix for the input of same axes of the swapaxes operator

src/operator/swapaxis-inl.h

@@ -35,6 +35,7 @@
 #include <string>
 #include <utility>
 #include "./operator_common.h"
+#include "./mshadow_op.h"
 
 namespace mxnet {
 namespace op {
@@ -63,7 +64,6 @@ template<typename xpu, typename DType>
 class SwapAxisOp : public Operator {
  public:
   explicit SwapAxisOp(SwapAxisParam p) {
-    CHECK_NE(p.dim1, p.dim2) << "dim1 can not be equal dim2.";
     this->param_ = p;
   }
 
@@ -131,6 +131,16 @@ class SwapAxisOp : public Operator {
 
     if (shape_in.Size() == 0U) return;
 
+    if (axis1 == axis2) {
+      if (out_req == kAddTo) {
+        mxnet_op::Kernel<mxnet_op::op_with_req<mshadow_op::identity, kAddTo>, xpu>::Launch(
+          s, data_out.Size(), data_out.dptr<DType>(), data_in.dptr<DType>());
+      } else {
+        mxnet_op::copy(s, data_out, data_in);
+      }
+      return;
+    }
+
     Shape<5> inter_shape;
 
     Reshape2Five(&inter_shape, shape_in, axis1, axis2);

tests/python/unittest/test_numpy_op.py

@@ -1967,9 +1967,11 @@ def check_comp_op(op_name, x1, x2):
 @with_seed()
 @use_np
 def test_np_swapaxes():
-    config = [((0, 1, 2), 0, 1),
-              ((0, 1, 2), -1, -2),
-              ((4, 5, 6, 7), 2, 3),
+    config = [((0, 1, 2), 0, 0),
+              ((0, 1, 2), 1, 2),
+              ((0, 1, 2), 1, -2),
+              ((4, 5, 6, 7), 1, 1),
+              ((4, 5, 6, 7), 2, -2),
               ((4, 5, 6, 7), -2, -3)]
 
     class TestSwapaxes(HybridBlock):

tests/python/unittest/test_operator.py

@@ -716,6 +716,23 @@ def test_swapaxes():
 
     assert_almost_equal(out, swap_)
 
+    config = [((1, 1, 2), 0, 1),
+              ((1, 1, 2), -1, -2),
+              ((4, 5, 6, 7), 1, 1),
+              ((4, 5, 6, 7), 2, 3),
+              ((4, 5, 6, 7), -2, 2),
+              ((4, 5, 6, 7), -2, -3)]
+
+    for shape, axis1, axis2 in config:
+        data_np = np.random.uniform(size=shape)
+        data_mx = mx.nd.array(data_np, dtype=data_np.dtype)
+        ret_np = np.swapaxes(data_np, axis1=axis1, axis2=axis2)
+        ret_mx = mx.symbol.SwapAxis(data, dim1=axis1, dim2=axis2)
+        exe_c = ret_mx.bind(default_context(), args=[data_mx])
+        exe_c.forward(is_train=True)
+        out = exe_c.outputs[0]
+        assert_almost_equal(out, ret_np)
+
 
 @with_seed()
 def test_scalarop():