[Python array API standard] Support Tensor.mT in eager/static mode (#68833)

HydrogenSulfate · SigureMo · web-flow · commit bb5d143847cf · 2024-10-23T13:07:50.000+08:00
* support Tensor.mT in eager/static mode

* update unittests

* fix

* fix DummyTracker usage

---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/base/dygraph/math_op_patch.py b/python/paddle/base/dygraph/math_op_patch.py
@@ -189,6 +189,17 @@ def _T_(var: Tensor) -> Tensor:
         out = _C_ops.transpose(var, perm)
         return out
 
+    @property
+    def _mT_(var: Tensor) -> Tensor:
+        if len(var.shape) < 2:
+            raise ValueError(
+                f"Tensor.ndim({var.ndim}) is required to be greater than or equal to 2."
+            )
+        perm = list(range(len(var.shape)))
+        perm[-1], perm[-2] = perm[-2], perm[-1]
+        out = _C_ops.transpose(var, perm)
+        return out
+
     eager_methods = [
         ('__neg__', _neg_),
         ('__abs__', _abs_),
@@ -203,6 +214,7 @@ def _T_(var: Tensor) -> Tensor:
         ('ndim', _ndim),
         ('size', _size_),
         ('T', _T_),
+        ('mT', _mT_),
         # for logical compare
         ('__array_ufunc__', None),
     ]
diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py
@@ -119,6 +119,7 @@ def _to_static_var(self, to_parameter=False, **kwargs):
         attr_not_need_keys = [
             'grad',
             'T',
+            'mT',
             'place',
             '_place_str',
             'data',
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -528,7 +528,24 @@ def T(self):
 
         perm = list(range(len(self.meta.shape) - 1, -1, -1))
         perm_var = ListVariable(perm, self.graph, tracker=ConstTracker(perm))
-        assert perm_var is not None
+        out = self.graph.call_paddle_api(paddle.transpose, self, perm_var)
+        return out
+
+    @tensor_property
+    def mT(self):
+        """
+        Return a new TensorVariable object that wraps the result of calling the mT method on the wrapped value of this TensorVariable.
+        """
+        from .container import ListVariable
+
+        if len(self.meta.shape) < 2:
+            raise ValueError(
+                f"Variable.ndim({self.ndim}) is required to be greater than or equal to 2."
+            )
+
+        perm = list(range(len(self.meta.shape)))
+        perm[-1], perm[-2] = perm[-2], perm[-1]
+        perm_var = ListVariable(perm, self.graph, tracker=DummyTracker([self]))
         out = self.graph.call_paddle_api(paddle.transpose, self, perm_var)
         return out
 
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -479,6 +479,39 @@ def _T_(self):
 
         return _C_ops.transpose(self, perm)
 
+    @property
+    def _mT_(self):
+        """
+
+        Permute current Value with its last two dimensions reversed.
+
+        If `n` is the dimensions of `x` , `x.mT` is equivalent to `x.transpose([0, 1, ..., n-1, n-2])`.
+
+        Examples:
+            .. code-block:: python
+
+                >>> import paddle
+                >>> paddle.enable_static()
+
+                >>> x = paddle.ones(shape=[2, 3, 5])
+                >>> x_mT = x.mT
+
+                >>> exe = paddle.static.Executor()
+                >>> x_mT_np = exe.run(paddle.static.default_main_program(), fetch_list=[x_mT])[0]
+                >>> print(x_mT_np.shape)
+                (2, 5, 3)
+
+        """
+        if len(self.shape) < 2:
+            raise ValueError(
+                f"Tensor.ndim({len(self.shape)}) is required to be greater than or equal to 2."
+            )
+
+        perm = list(range(len(self.shape)))
+        perm[-1], perm[-2] = perm[-2], perm[-1]
+
+        return _C_ops.transpose(self, perm)
+
     def _int_(self):
         error_msg = """\
             int(Tensor) is not supported in static graph mode. Because it's value is not available during the static mode.
@@ -973,6 +1006,7 @@ def register_hook(self, hook):
         ('astype', astype),
         ('size', _size_),
         ('T', _T_),
+        ('mT', _mT_),
         ('clone', clone),
         ('clear_gradient', clear_gradient),
         ('append', append),
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -464,6 +464,69 @@ def test_T(self):
                     (output_x,) = exe.run(main_program, fetch_list=[x_T])
                     self.assertEqual(output_x.shape, tuple(out_shape))
 
+    def test_mT(self):
+        with paddle.pir_utils.IrGuard():
+            shape = [1]
+            x = paddle.rand(shape, dtype="float32")
+            self.assertRaises(ValueError, getattr, x, 'mT')
+
+            for ndim in range(2, 5):
+                # shape is [1, 2], [1, 2, 3], [1, 2, 3, 4]
+                shape = list(range(1, ndim + 1))
+                out_shape = list(shape)
+                out_shape[-2], out_shape[-1] = out_shape[-1], out_shape[-2]
+                main_program, exe, program_guard = new_program()
+                with program_guard:
+                    x = paddle.rand(shape, dtype="float32")
+                    x_mT = x.mT
+                    self.assertEqual(x_mT.shape, out_shape)
+                    (output_x,) = exe.run(main_program, fetch_list=[x_mT])
+                    self.assertEqual(output_x.shape, tuple(out_shape))
+
+            shape = [1, 2, 3, 0, 1]
+            out_shape = list(shape)
+            out_shape[-2], out_shape[-1] = out_shape[-1], out_shape[-2]
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.rand(shape, dtype="float32")
+                x_mT = x.mT
+                self.assertEqual(x_mT.shape, out_shape)
+                (output_x,) = exe.run(main_program, fetch_list=[x_mT])
+                self.assertEqual(output_x.shape, tuple(out_shape))
+
+            shape = [1, 2, 3, 1, 0]
+            out_shape = list(shape)
+            out_shape[-2], out_shape[-1] = out_shape[-1], out_shape[-2]
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.rand(shape, dtype="float32")
+                x_mT = x.mT
+                self.assertEqual(x_mT.shape, out_shape)
+                (output_x,) = exe.run(main_program, fetch_list=[x_mT])
+                self.assertEqual(output_x.shape, tuple(out_shape))
+
+            shape = [1, 2, 3, 0, 0]
+            out_shape = list(shape)
+            out_shape[-2], out_shape[-1] = out_shape[-1], out_shape[-2]
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.rand(shape, dtype="float32")
+                x_mT = x.mT
+                self.assertEqual(x_mT.shape, out_shape)
+                (output_x,) = exe.run(main_program, fetch_list=[x_mT])
+                self.assertEqual(output_x.shape, tuple(out_shape))
+
+            shape = [0, 2, 3, 0, 0]
+            out_shape = list(shape)
+            out_shape[-2], out_shape[-1] = out_shape[-1], out_shape[-2]
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.rand(shape, dtype="float32")
+                x_mT = x.mT
+                self.assertEqual(x_mT.shape, out_shape)
+                (output_x,) = exe.run(main_program, fetch_list=[x_mT])
+                self.assertEqual(output_x.shape, tuple(out_shape))
+
     def test_hash(self):
         with paddle.pir_utils.IrGuard():
             _, _, program_guard = new_program()
diff --git a/test/legacy_test/test_math_op_patch_var_base.py b/test/legacy_test/test_math_op_patch_var_base.py
@@ -825,6 +825,18 @@ def test_tensor_patch_method(self):
         self.assertTrue(x_T.shape, [7, 9, 6, 3])
         np.testing.assert_array_equal(x_T.numpy(), x_np.T)
 
+        x_np = np.random.randn(3, 6, 9, 7)
+        x = paddle.to_tensor(x_np)
+        x_mT = x.mT
+        self.assertTrue(x_mT.shape, [3, 6, 7, 9])
+        np.testing.assert_array_equal(
+            x_mT.numpy(), x_np.transpose([0, 1, 3, 2])
+        )
+
+        x_np = np.random.randn(3)
+        x = paddle.to_tensor(x_np)
+        self.assertRaises(ValueError, getattr, x, "mT")
+
         self.assertTrue(inspect.ismethod(a.dot))
         self.assertTrue(inspect.ismethod(a.logsumexp))
         self.assertTrue(inspect.ismethod(a.multiplex))
diff --git a/test/sot/test_18_tensor_method.py b/test/sot/test_18_tensor_method.py
@@ -57,6 +57,10 @@ def tensor_method_property(a: paddle.Tensor, b: paddle.Tensor):
     )
 
 
+def tensor_method_property_mT(a: paddle.Tensor):
+    return a.mT
+
+
 def middle_tensor_name(a: paddle.Tensor, b: paddle.Tensor):
     c = a + b
     return c.name
@@ -90,6 +94,12 @@ def test_middle_tensor_name(self):
         y = paddle.rand([42, 24])
         self.assert_results(middle_tensor_name, x, y)
 
+    def test_tensor_method_property_mT(self):
+        x = paddle.rand([42, 24, 2, 2, 3, 2], dtype='float64')
+        y = paddle.rand([42, 24, 2, 3, 3, 2], dtype='float32')
+        self.assert_results(tensor_method_property_mT, x)
+        self.assert_results(tensor_method_property_mT, y)
+
 
 if __name__ == "__main__":
     unittest.main()
