[Paddle Tensor No.6][BUPT] 新增 Tensor.__ror__ (#69463)

ZHOU05030 · web-flow · commit f2cd9aacdf6e · 2024-11-25T21:27:04.000+08:00
* add tensor.ror

* 修改测试文件

* 修改测试文件

* 修改测试文件

* 修改单测文件

* 修改init文件

* 添加 __ror__

* 添加动态图测试

* 修改bug

* 添加单侧动态测试

* 修改动态图测试

* 增加数据类型判断

* 添加test_math_op_patch文件

* 修改bug

* 修改静态图测试

* 添加pir的测试部分

* 修改if逻辑

* 修改pir测试

* 添加tensor.prototype.pyi

* Resolve merge conflict in __init__.py

* Resolve merge conflict in __init__.py

* Resolve merge conflict in tensor.__init__.py

* pre-commit
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
@@ -105,6 +105,7 @@
 )
 from .logic import (  # noqa: F401
     __rand__,
+    __ror__,
     allclose,
     bitwise_and,
     bitwise_and_,
@@ -867,6 +868,7 @@
     ('__and__', 'bitwise_and'),
     ('__rand__', '__rand__'),
     ('__or__', 'bitwise_or'),
+    ('__ror__', '__ror__'),
     ('__xor__', 'bitwise_xor'),
     ('__invert__', 'bitwise_not'),
     ('__pos__', 'positive'),
diff --git a/python/paddle/tensor/logic.py b/python/paddle/tensor/logic.py
@@ -1343,6 +1343,21 @@ def bitwise_or(
     )
 
 
+def __ror__(
+    x: Tensor,
+    y: int | bool,
+    out: Tensor | None = None,
+    name: str | None = None,
+) -> Tensor:
+    if isinstance(y, (int, bool)):
+        y = paddle.to_tensor(y, dtype=x.dtype)
+        return bitwise_or(y, x, out=out, name=name)
+    else:
+        raise TypeError(
+            f"unsupported operand type(s) for |: '{type(y).__name__}' and 'Tensor'"
+        )
+
+
 @inplace_apis_in_dygraph_only
 def bitwise_or_(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
     r"""
diff --git a/python/paddle/tensor/tensor.prototype.pyi b/python/paddle/tensor/tensor.prototype.pyi
@@ -171,6 +171,7 @@ class AbstractTensor:
     def __mod__(self, y: _typing.TensorLike) -> Tensor: ...
     def __pow__(self, y: _typing.TensorLike) -> Tensor: ...
     def __and__(self, y: _typing.TensorLike) -> Tensor: ...
+    def __ror__(self, y: _typing.TensorLike) -> Tensor: ...
     def __div__(self, y: _typing.TensorLike) -> Tensor: ...
     def __radd__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rsub__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -335,6 +335,42 @@ def test_bitwise_or(self):
         )
         np.testing.assert_array_equal(out[0], out_np)
 
+    @prog_scope()
+    def test_ror(self):
+        place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+        x_int = 5
+        y_np = np.random.randint(-100, 100, [2, 3, 5]).astype("int32")
+        y = paddle.static.data("y", y_np.shape, dtype=y_np.dtype)
+        z = x_int | y
+        exe = paddle.static.Executor(place)
+        out = exe.run(
+            feed={'y': y_np},
+            fetch_list=[z],
+        )
+        out_ref = x_int | y_np
+        np.testing.assert_array_equal(out[0], out_ref)
+        x_bool = True
+        res_ror_bool = x_bool | y
+        out_bool = exe.run(
+            feed={'y': y_np},
+            fetch_list=[res_ror_bool],
+        )
+        res_py_bool = x_bool | y_np
+        np.testing.assert_array_equal(out_bool[0], res_py_bool)
+
+        for x_invalid in (
+            np.float32(5.0),
+            np.float64(5.0),
+            np.complex64(5),
+            np.complex128(5.0 + 2j),
+        ):
+            with self.assertRaises(TypeError):
+                x_invalid | y
+
     @prog_scope()
     def test_bitwise_xor(self):
         x_np = np.random.randint(-100, 100, [2, 3, 5]).astype("int32")
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -17,6 +17,7 @@
 import warnings
 
 import numpy as np
+from utils import dygraph_guard, static_guard
 
 import paddle
 from paddle import base
@@ -251,6 +252,63 @@ def test_bitwise_or(self):
                 np.testing.assert_array_equal(res_np_c, c_np)
                 np.testing.assert_array_equal(res_np_d, d_np)
 
+    def test_ror(self):
+        with dygraph_guard():
+            x_int32 = 5
+            x_bool = True
+            y_np = np.random.randint(0, 2, [2, 3, 5]).astype("int32")
+            y_tensor = paddle.to_tensor(y_np)
+            res_ror_int32 = x_int32 | y_tensor
+            res_py_int32 = x_int32 | y_tensor.numpy()
+            np.testing.assert_array_equal(res_py_int32, res_ror_int32.numpy())
+            res_ror_bool = x_bool | y_tensor
+            res_py_bool = x_bool | y_tensor.numpy()
+            np.testing.assert_array_equal(res_py_bool, res_ror_bool.numpy())
+            for x_np in (
+                np.float32(5.0),
+                np.float64(5.0),
+                np.complex64(5),
+                np.complex128(5.0 + 2j),
+            ):
+                with self.assertRaises(TypeError):
+                    x_np | y_tensor
+
+        with static_guard():
+            with paddle.pir_utils.IrGuard():
+                main_program, exe, program_guard = new_program()
+                with program_guard:
+                    x_int = 5
+                    y_np = np.random.randint(-100, 100, [2, 3, 5]).astype(
+                        "int32"
+                    )
+                    y = paddle.static.data("y", y_np.shape, dtype=y_np.dtype)
+                    z = x_int | y
+                    out = exe.run(
+                        main_program,
+                        feed={'y': y_np},
+                        fetch_list=[z],
+                    )
+                    out_ref = x_int | y_np
+                    np.testing.assert_array_equal(out[0], out_ref)
+                    x_bool = True
+                    res_ror_bool = x_bool | y
+                    out_bool = exe.run(
+                        main_program,
+                        feed={'y': y_np},
+                        fetch_list=[res_ror_bool],
+                    )
+                    res_py_bool = x_bool | y_np
+                    np.testing.assert_array_equal(out_bool[0], res_py_bool)
+
+                    for x_invalid in (
+                        np.float32(5.0),
+                        np.float64(5.0),
+                        np.complex64(5),
+                        np.complex128(5.0 + 2j),
+                    ):
+                        with self.assertRaises(TypeError):
+                            x_invalid | y
+
     def test_bitwise_and(self):
         paddle.disable_static()
         x_np = np.random.randint(-100, 100, [2, 3, 5]).astype("int32")
