[Zero-Dim] support input 0D for paddle.moveaxis / quantile (#49813)

* [Zero-Dim] support input 0D for paddle.moveaxis/quantile

* fix CI

paddle/phi/kernels/xpu/transpose_grad_kernel.cc

@@ -26,6 +26,14 @@ void TransposeGradKernel(const Context& dev_ctx,
                          DenseTensor* x_grad) {
   using XPUType = typename XPUTypeTrait<T>::Type;
   dev_ctx.template Alloc<T>(x_grad);
+  if (x_grad->numel() == 0) {
+    return;
+  }
+  if (axis.size() == 0) {
+    phi::Copy<Context>(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    return;
+  }
+
   std::vector<int> reversed_axis(axis);
   for (size_t i = 0; i < axis.size(); i++) {
     reversed_axis[axis[i]] = i;

python/paddle/fluid/tests/unittests/test_quantile_and_nanquantile.py

@@ -209,12 +209,6 @@ class TestError(unittest.TestCase):
 
         self.assertRaises(ValueError, test_axis_value_error_2)
 
-        # Test error with no valid axis
-        def test_axis_value_error_3():
-            paddle_res = paddle.quantile(self.x, q=0.4, axis=[])
-
-        self.assertRaises(ValueError, test_axis_value_error_3)
-
 
 class TestQuantileRuntime(unittest.TestCase):
     """

python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py

@@ -18,7 +18,6 @@ import numpy as np
 from decorator_helper import prog_scope
 
 import paddle
-import paddle.fluid as fluid
 import paddle.nn.functional as F
 
 unary_api_list = [
@@ -100,8 +99,8 @@ class TestUnaryAPI(unittest.TestCase):
         for api in unary_api_list:
             x = paddle.rand([])
             x.stop_gradient = False
-            x.retain_grads()
             out = api(x)
+
             out.retain_grads()
             out.backward()
 
@@ -123,10 +122,12 @@ class TestUnaryAPI(unittest.TestCase):
         paddle.enable_static()
 
         for api in unary_api_list:
-            main_prog = fluid.Program()
+            main_prog = paddle.static.Program()
             block = main_prog.global_block()
             exe = paddle.static.Executor()
-            with fluid.program_guard(main_prog, fluid.Program()):
+            with paddle.static.program_guard(
+                main_prog, paddle.static.Program()
+            ):
                 x = paddle.rand([])
                 x.stop_gradient = False
                 out = api(x)
@@ -202,29 +203,34 @@ class TestReduceAPI(unittest.TestCase):
             else:
                 x = paddle.rand([])
             x.stop_gradient = False
-            x.retain_grads()
             out = api(x, None)
+
             out.retain_grads()
             out.backward()
 
+            out_empty_list = api(x, [])
+            self.assertEqual(out_empty_list, out)
+
             self.assertEqual(x.shape, [])
             self.assertEqual(out.shape, [])
-            self.assertEqual(out.numpy(), x.numpy())
+            np.testing.assert_allclose(out.numpy(), x.numpy())
             if x.grad is not None:
                 self.assertEqual(x.grad.shape, [])
-                self.assertEqual(x.grad.numpy(), 1.0)
                 self.assertEqual(out.grad.shape, [])
-                self.assertEqual(out.grad.numpy(), 1.0)
+                np.testing.assert_allclose(x.grad.numpy(), np.array(1.0))
+                np.testing.assert_allclose(out.grad.numpy(), np.array(1.0))
 
         paddle.enable_static()
 
     def test_static_reduce(self):
         paddle.enable_static()
         for api in reduce_api_list:
-            main_prog = fluid.Program()
+            main_prog = paddle.static.Program()
             block = main_prog.global_block()
             exe = paddle.static.Executor()
-            with fluid.program_guard(main_prog, fluid.Program()):
+            with paddle.static.program_guard(
+                main_prog, paddle.static.Program()
+            ):
                 # 1) x is 0D
                 if api in [paddle.all, paddle.any]:
                     x = paddle.randint(0, 2, []).astype('bool')
@@ -234,6 +240,9 @@ class TestReduceAPI(unittest.TestCase):
                 out = api(x, None)
                 paddle.static.append_backward(out.sum())
 
+                out_empty_list = api(x, None)
+                self.assertEqual(out_empty_list.shape, ())
+
                 fetch_list = [x, out]
                 if block.has_var(x.grad_name):
                     fetch_list.extend([x.grad_name, out.grad_name])
@@ -241,12 +250,12 @@ class TestReduceAPI(unittest.TestCase):
                 res = exe.run(main_prog, fetch_list=fetch_list)
                 self.assertEqual(res[0].shape, ())
                 self.assertEqual(res[1].shape, ())
-                self.assertEqual(res[0], res[1])
+                np.testing.assert_allclose(res[0], res[1])
                 if len(res) > 2:
                     self.assertEqual(res[2].shape, ())
                     self.assertEqual(res[3].shape, ())
-                    self.assertEqual(res[2], 1.0)
-                    self.assertEqual(res[3], 1.0)
+                    np.testing.assert_allclose(res[2], np.array(1.0))
+                    np.testing.assert_allclose(res[3], np.array(1.0))
 
         paddle.disable_static()
 
@@ -293,8 +302,6 @@ class TestBinaryAPI(unittest.TestCase):
             y = paddle.rand([])
             x.stop_gradient = False
             y.stop_gradient = False
-            x.retain_grads()
-            y.retain_grads()
             if isinstance(api, dict):
                 out = api['func'](x, y)
                 out_cls = getattr(paddle.Tensor, api['cls_method'])(x, y)
@@ -318,8 +325,6 @@ class TestBinaryAPI(unittest.TestCase):
             y = paddle.rand([])
             x.stop_gradient = False
             y.stop_gradient = False
-            x.retain_grads()
-            y.retain_grads()
             if isinstance(api, dict):
                 out = api['func'](x, y)
                 out_cls = getattr(paddle.Tensor, api['cls_method'])(x, y)
@@ -341,8 +346,6 @@ class TestBinaryAPI(unittest.TestCase):
             # 3) x is 0D , y is ND
             x = paddle.rand([])
             y = paddle.rand([2, 3, 4])
-            x.retain_grads()
-            y.retain_grads()
             x.stop_gradient = False
             y.stop_gradient = False
             if isinstance(api, dict):
@@ -366,10 +369,10 @@ class TestBinaryAPI(unittest.TestCase):
             # 4) x is 0D , y is scalar
             x = paddle.rand([])
             x.stop_gradient = False
-            x.retain_grads()
             y = 0.5
             if isinstance(api, dict):
                 out = getattr(paddle.Tensor, api['cls_method'])(x, y)
+
                 out.retain_grads()
                 out.backward()
 
@@ -403,9 +406,11 @@ class TestBinaryAPI(unittest.TestCase):
     def test_static_binary(self):
         paddle.enable_static()
         for api in binary_api_list:
-            main_prog = fluid.Program()
+            main_prog = paddle.static.Program()
             block = main_prog.global_block()
-            with fluid.program_guard(main_prog, fluid.Program()):
+            with paddle.static.program_guard(
+                main_prog, paddle.static.Program()
+            ):
                 # 1) x is 0D, y is 0D
                 x = paddle.rand([])
                 y = paddle.rand([])
@@ -511,8 +516,10 @@ class TestBinaryAPI(unittest.TestCase):
                 '''
 
         for api in binary_int_api_list:
-            main_prog = fluid.Program()
-            with fluid.program_guard(main_prog, fluid.Program()):
+            main_prog = paddle.static.Program()
+            with paddle.static.program_guard(
+                main_prog, paddle.static.Program()
+            ):
                 # 1) x is 0D, y is 0D
                 x = paddle.randint(-10, 10, [])
                 y = paddle.randint(-10, 10, [])
@@ -541,10 +548,43 @@ class TestSundryAPI(unittest.TestCase):
         paddle.disable_static()
         self.x = paddle.rand([])
 
+    def test_quantile(self):
+        # 1) x is 0D
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.quantile(x, 0.5, axis=None)
+
+        out.retain_grads()
+        out.backward()
+
+        out_empty_list = paddle.quantile(x, 0.5, axis=[])
+        self.assertEqual(out_empty_list, out)
+
+        self.assertEqual(x.shape, [])
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out, x)
+
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad, 1.0)
+        self.assertEqual(out.grad.shape, [])
+        self.assertEqual(out.grad, 1.0)
+
+        # 2) x is ND
+        x = paddle.rand([2, 3])
+        x.stop_gradient = False
+        out = paddle.quantile(x, 0.5, axis=None)
+
+        out.retain_grads()
+        out.backward()
+
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out.grad.shape, [])
+        self.assertEqual(out.grad, 1.0)
+        self.assertEqual(x.grad.shape, [2, 3])
+
     def test_flip(self):
         x = paddle.rand([])
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.flip(x, axis=[])
         out.retain_grads()
         out.backward()
@@ -636,7 +676,6 @@ class TestSundryAPI(unittest.TestCase):
     def test_pow_factor(self):
         x = paddle.rand([])
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.pow(x, 2.0)
         out.retain_grads()
         out.backward()
@@ -648,7 +687,6 @@ class TestSundryAPI(unittest.TestCase):
     def test_cast(self):
         x = paddle.full([], 1.0, 'float32')
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.cast(x, 'int32')
         out.retain_grads()
         out.backward()
@@ -660,7 +698,6 @@ class TestSundryAPI(unittest.TestCase):
     def test_cumprod(self):
         x = paddle.full([], 1.0, 'float32')
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.cumprod(x, 0)
         out.retain_grads()
         out.backward()
@@ -675,7 +712,6 @@ class TestSundryAPI(unittest.TestCase):
     def test_clip(self):
         x = paddle.uniform([], None, -10, 10)
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.clip(x, -5, 5)
         out.retain_grads()
         out.backward()
@@ -687,7 +723,6 @@ class TestSundryAPI(unittest.TestCase):
     def test_increment(self):
         x = paddle.rand([])
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.increment(x, 1.0)
         out.retain_grads()
         out.backward()
@@ -711,18 +746,49 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(out.shape, [])
 
     def test_searchsorted(self):
+        # have no backward
         x = paddle.to_tensor([1, 3, 5, 7, 9])
         y = paddle.rand([])
 
-        # only has forward kernel
         out = paddle.searchsorted(x, y)
 
         self.assertEqual(out.shape, [])
         self.assertEqual(out.numpy(), 0)
 
+    def test_transpose(self):
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.transpose(x, [])
+        out.retain_grads()
+        out.backward()
+
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out, x)
+        self.assertEqual(out.grad.shape, [])
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad, 1.0)
+
+        with self.assertRaises(ValueError):
+            x = paddle.transpose(x, [0])
+
+    def test_moveaxis(self):
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.moveaxis(x, [], [])
+        out.retain_grads()
+        out.backward()
+
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out, x)
+        self.assertEqual(out.grad.shape, [])
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad, 1.0)
+
+        with self.assertRaises(AssertionError):
+            x = paddle.moveaxis(x, [1], [0])
+
     def test_gather_1D(self):
         x = paddle.to_tensor([1.0, 3.0, 5.0, 7.0, 9.0], stop_gradient=False)
-        x.retain_grads()
         index = paddle.full([], 2, 'int64')
         out = paddle.gather(x, index)
         out.retain_grads()
@@ -737,7 +803,6 @@ class TestSundryAPI(unittest.TestCase):
         x = paddle.to_tensor(
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], stop_gradient=False
         )
-        x.retain_grads()
         index = paddle.full([], 1, 'int64')
         out = paddle.gather(x, index)
         out.retain_grads()
@@ -752,7 +817,6 @@ class TestSundryAPI(unittest.TestCase):
         x = paddle.to_tensor(
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], stop_gradient=False
         )
-        x.retain_grads()
         index = paddle.full([], 1, 'int64')
         out = paddle.gather(x, index, axis=1)
         out.retain_grads()
@@ -763,9 +827,8 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(x.grad.shape, [2, 3])
         self.assertEqual(out.grad.shape, [2])
 
-    def test_scatter_1D(self):
+    def _test_scatter_1D(self):
         x = paddle.to_tensor([1.0, 3.0, 5.0, 7.0, 9.0], stop_gradient=False)
-        x.retain_grads()
         index = paddle.full([], 2, 'int64')
         updates = paddle.full([], 4.0)
         out = paddle.scatter(x, index, updates)
@@ -776,7 +839,7 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(out.numpy()[2], 4)
         self.assertEqual(out.grad.shape, [5])
 
-    def test_scatter_XD(self):
+    def _test_scatter_XD(self):
         x = paddle.to_tensor(
             [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], stop_gradient=False
         )
@@ -866,7 +929,6 @@ class TestSundryAPI(unittest.TestCase):
 
             x = paddle.randn(())
             x.stop_gradient = False
-            x.retain_grads()
 
             out = paddle.kthvalue(x, 1)
             out[0].backward()
@@ -887,7 +949,6 @@ class TestSundryAPI(unittest.TestCase):
             paddle.set_device(place)
 
             x = paddle.randn(())
-            x.retain_grads()
             x.stop_gradient = False
 
             out = paddle.mode(x)
@@ -904,7 +965,6 @@ class TestSundryAPI(unittest.TestCase):
     def test_flatten(self):
         x = paddle.rand([])
         x.stop_gradient = False
-        x.retain_grads()
 
         start_axis = 0
         stop_axis = -1
@@ -925,8 +985,8 @@ class TestSundryAPI(unittest.TestCase):
     def test_scale(self):
         x = paddle.rand([])
         x.stop_gradient = False
-        x.retain_grads()
         out = paddle.scale(x, scale=2.0, bias=1.0)
+
         out.retain_grads()
         out.backward()
 
@@ -1018,9 +1078,8 @@ class TestSundryAPI(unittest.TestCase):
     def test_reshape_list(self):
         x = paddle.rand([])
         x.stop_gradient = False
-        x.retain_grads()
-
         out = paddle.reshape(x, [])
+
         out.retain_grads()
         out.backward()
         self.assertEqual(x.grad.shape, [])
@@ -1050,10 +1109,9 @@ class TestSundryAPI(unittest.TestCase):
 
     def test_reshape_tensor(self):
         x = paddle.rand([1, 1])
-        x.retain_grads()
         x.stop_gradient = False
-
         out = paddle.reshape(x, [])
+
         out.retain_grads()
         out.backward()
         self.assertEqual(x.grad.shape, [1, 1])
@@ -1237,7 +1295,6 @@ class TestSundryAPI(unittest.TestCase):
 
             x = paddle.randn(())
             x.stop_gradient = False
-            x.retain_grads()
 
             out = paddle.repeat_interleave(x, 2, None)
             out.backward()
@@ -1379,6 +1436,44 @@ class TestSundryAPIStatic(unittest.TestCase):
         paddle.enable_static()
         self.exe = paddle.static.Executor()
 
+    @prog_scope()
+    def test_quantile(self):
+        x1 = paddle.rand([])
+        x1.stop_gradient = False
+        out1 = paddle.quantile(x1, 0.5, axis=None)
+        paddle.static.append_backward(out1.sum())
+
+        x2 = paddle.rand([2, 3])
+        x2.stop_gradient = False
+        out2 = paddle.quantile(x2, 0.5, axis=None)
+        paddle.static.append_backward(out2.sum())
+
+        out_empty_list = paddle.quantile(x1, 0.5, axis=[])
+        self.assertEqual(out_empty_list.shape, ())
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(
+            prog,
+            fetch_list=[
+                out1,
+                out2,
+                x1.grad_name,
+                out1.grad_name,
+                x2.grad_name,
+                out2.grad_name,
+            ],
+        )
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[1].shape, ())
+
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[3].shape, ())
+        self.assertEqual(res[3], 1.0)
+
+        self.assertEqual(res[4].shape, (2, 3))
+        self.assertEqual(res[5].shape, ())
+        self.assertEqual(res[5], 1.0)
+
     @prog_scope()
     def test_flip(self):
         x = paddle.rand([])
@@ -1509,6 +1604,42 @@ class TestSundryAPIStatic(unittest.TestCase):
         self.assertEqual(res[0].shape, ())
         self.assertEqual(res[0], 0)
 
+    @prog_scope()
+    def test_transpose(self):
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.transpose(x, [])
+        paddle.static.append_backward(out.sum())
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, fetch_list=[out, x.grad_name, out.grad_name])
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+
+        with self.assertRaises(ValueError):
+            x = paddle.transpose(x, [0])
+
+    @prog_scope()
+    def test_moveaxis(self):
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.moveaxis(x, [], [])
+        paddle.static.append_backward(out.sum())
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, fetch_list=[out, x.grad_name, out.grad_name])
+        self.assertEqual(res[0].shape, ())
+        self.assertEqual(res[1].shape, ())
+        self.assertEqual(res[1], 1.0)
+        self.assertEqual(res[2].shape, ())
+        self.assertEqual(res[2], 1.0)
+
+        with self.assertRaises(AssertionError):
+            x = paddle.moveaxis(x, [0], [1])
+
     @prog_scope()
     def test_gather_1D(self):
         x = paddle.full([10], 1.0, 'float32')
@@ -2487,7 +2618,7 @@ class TestNoBackwardAPIStatic(unittest.TestCase):
         ids = paddle.full(shape=[], fill_value=1, dtype='int64')
         emb = paddle.static.nn.embedding(ids, (20, 3))
         prog = paddle.static.default_main_program()
-        self.exe.run(paddle.fluid.default_startup_program())
+        self.exe.run(paddle.static.default_startup_program())
         res = self.exe.run(prog, fetch_list=[emb])
         self.assertEqual(res[0].shape, (3,))
 
@@ -2495,7 +2626,7 @@ class TestNoBackwardAPIStatic(unittest.TestCase):
         label = paddle.full(shape=[], fill_value=2, dtype='int64')
         one_hot_label = paddle.nn.functional.one_hot(label, num_classes=4)
         prog = paddle.static.default_main_program()
-        self.exe.run(paddle.fluid.default_startup_program())
+        self.exe.run(paddle.static.default_startup_program())
         res = self.exe.run(prog, fetch_list=[one_hot_label])
 
         self.assertEqual(res[0].shape, (4,))

python/paddle/fluid/tests/unittests/xpu/test_zero_dim_tensor_xpu.py

@@ -17,12 +17,10 @@ import unittest
 import numpy as np
 
 import paddle
-import paddle.fluid as fluid
 import paddle.nn.functional as F
 
 paddle.set_device('xpu')
-fluid.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
-
+paddle.set_flags({"FLAGS_retain_grad_for_all_tensor": True})
 
 unary_api_list = [
     paddle.nn.functional.elu,
@@ -149,13 +147,17 @@ class TestReduceAPI(unittest.TestCase):
                 x = paddle.rand([])
             x.stop_gradient = False
             out = api(x, None)
+
             out.backward()
 
             self.assertEqual(x.shape, [])
             self.assertEqual(out.shape, [])
+            np.testing.assert_allclose(out.numpy(), x.numpy())
             if x.grad is not None:
                 self.assertEqual(x.grad.shape, [])
                 self.assertEqual(out.grad.shape, [])
+                np.testing.assert_allclose(x.grad.numpy(), np.array(1.0))
+                np.testing.assert_allclose(out.grad.numpy(), np.array(1.0))
 
         paddle.enable_static()
 
@@ -440,6 +442,36 @@ class TestSundryAPI(unittest.TestCase):
         self.assertEqual(out.shape, [])
         self.assertEqual(out.numpy(), 0)
 
+    def test_transpose(self):
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.transpose(x, [])
+        out.backward()
+
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out, x)
+        self.assertEqual(out.grad.shape, [])
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad, 1.0)
+
+        with self.assertRaises(ValueError):
+            x = paddle.transpose(x, [0])
+
+    def test_moveaxis(self):
+        x = paddle.rand([])
+        x.stop_gradient = False
+        out = paddle.moveaxis(x, [], [])
+        out.backward()
+
+        self.assertEqual(out.shape, [])
+        self.assertEqual(out, x)
+        self.assertEqual(out.grad.shape, [])
+        self.assertEqual(x.grad.shape, [])
+        self.assertEqual(x.grad, 1.0)
+
+        with self.assertRaises(AssertionError):
+            x = paddle.moveaxis(x, [1], [0])
+
     def test_gather_1D(self):
         x = paddle.to_tensor([1.0, 3.0, 5.0, 7.0, 9.0], stop_gradient=False)
         index = paddle.full([], 2, 'int64')

python/paddle/tensor/manipulation.py

@@ -14,8 +14,6 @@
 
 # TODO: define functions to manipulate a tensor
 
-from collections import Counter
-
 import numpy as np
 
 import paddle
@@ -4335,11 +4333,9 @@ def moveaxis(x, source, destination, name=None):
         dst
     ), "'source' must have the same number with 'destination'"
 
-    count = Counter(src).most_common(1)
-    if count[0][1] > 1:
+    if len(src) != len(set(src)):
         raise ValueError("Each elemment of 'source' must be unique!")
-    count = Counter(dst).most_common(1)
-    if count[0][1] > 1:
+    if len(dst) != len(set(dst)):
         raise ValueError("Each elemment of 'destination' must be unique!")
 
     ndim = len(x.shape)

python/paddle/tensor/stat.py

@@ -500,8 +500,6 @@ def _compute_quantile(x, q, axis=None, keepdim=False, ignore_nan=False):
         out_shape = [1] * dims
     else:
         if isinstance(axis, list):
-            if len(axis) <= 0:
-                raise ValueError("axis should not be empty")
             axis_src, axis_dst = [], []
             for axis_single in axis:
                 if not isinstance(axis_single, int) or not (
@@ -514,10 +512,15 @@ def _compute_quantile(x, q, axis=None, keepdim=False, ignore_nan=False):
                     axis_single = axis_single + dims
                 axis_src.append(axis_single)
                 out_shape[axis_single] = 1
+
             axis_dst = list(range(-len(axis), 0))
             x = paddle.moveaxis(x, axis_src, axis_dst)
-            x = paddle.flatten(x, axis_dst[0], axis_dst[-1])
-            axis = axis_dst[0]
+            if len(axis_dst) == 0:
+                x = paddle.flatten(x)
+                axis = 0
+            else:
+                x = paddle.flatten(x, axis_dst[0], axis_dst[-1])
+                axis = axis_dst[0]
         else:
             if not isinstance(axis, int) or not (axis < dims and axis >= -dims):
                 raise ValueError(