【SCU】【Paddle Tensor No.3】新增 Tensor.__rmod__，复用已有接口 Tensor.__mod__ (#69231)

PolaKuma · web-flow · commit d9c4b9f90f95 · 2024-11-12T10:31:18.000+08:00
* Add legend for masked_fill API

* add_rmod

* Update manipulation.py

* fix codestyle

* update test_dygraph_rmod

* Update the annotations
diff --git a/paddle/fluid/pybind/eager_math_op_patch.cc b/paddle/fluid/pybind/eager_math_op_patch.cc
@@ -135,7 +135,18 @@ std::set<phi::DataType> _complex_dtypes{
 //     '__truediv__',
 //     '__rdiv__',
 //     '__rtruediv__',
+//     '__floordiv__',
+//     '__pow__',
+//     '__rpow__',
+//     '__mod__',
+//     '__rmod__',
 //     '__matmul__',
+//     '__gt__',
+//     '__ge__',
+//     '__lt__',
+//     '__le__',
+//     '__eq__',
+//     '__ne__'
 
 void SetDevice(phi::Place place) {
   if (phi::is_gpu_place(place)) {
@@ -1059,6 +1070,97 @@ static PyObject* tensor__mod__method(TensorObject* self,
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+static PyObject* tensor__rmod__method(TensorObject* self,
+                                      PyObject* args,
+                                      PyObject* kwargs) {
+  phi::RecordEvent pythonc_record_event(
+      "__rmod__ pybind_patch_func", phi::TracerEventType::UserDefined, 1);
+  EAGER_TRY
+
+  VLOG(6) << "Running Eager tensor__rmod__method";
+
+  // Set Device ID
+  auto place = egr::Controller::Instance().GetExpectedPlace();
+  SetDevice(place);
+
+  paddle::Tensor ret;
+
+  paddle::Tensor self_tensor = self->tensor;
+  PyObject* other_obj = PyTuple_GET_ITEM(args, 0);
+
+  // 1. scalar exists cases
+  // there is no scalar_mod function for __rmod__ now
+  if (PyFloat_Check(other_obj) || PyCheckInteger(other_obj) ||
+      IsNumpyType(other_obj)) {
+    if (PyFloat_Check(other_obj)) {
+      if (_supported_int_dtype_.find(self_tensor.dtype()) !=
+          _supported_int_dtype_.end()) {
+        eager_gil_scoped_release guard;
+        self_tensor = cast_ad_func(self_tensor, DataType::FLOAT32);
+      }
+    } else if (PyCheckInteger(other_obj) &&
+               self_tensor.dtype() == DataType::BOOL) {
+      eager_gil_scoped_release guard;
+      self_tensor = cast_ad_func(self_tensor, DataType::INT64);
+    }
+  } else if (PyComplex_Check(other_obj)) {
+    if (is_support_complex(self_tensor.dtype()) == false) {
+      eager_gil_scoped_release guard;
+      self_tensor = cast_ad_func(
+          self_tensor, promoteTypes(self_tensor.dtype(), DataType::COMPLEX64));
+    }
+  }
+
+  // 2. create or get tensor for other_obj
+  paddle::Tensor other_tensor;
+  if (PyCheckTensor(other_obj)) {
+    auto& self_tensor_ref_addr = self->tensor;
+    auto& other_tensor_ref_addr = CastPyArg2Tensor(other_obj, 0);
+    const phi::distributed::ProcessMesh* mesh = nullptr;
+    if (InputsContainDistTensor(
+            &mesh, self_tensor_ref_addr, other_tensor_ref_addr)) {
+      ConvertAllInputsToDistTensor(
+          mesh, self_tensor_ref_addr, other_tensor_ref_addr);
+    }
+    self_tensor = self_tensor_ref_addr;
+    other_tensor = other_tensor_ref_addr;
+  } else {
+    if (IsNumpyArray(other_obj)) {
+      py::object numpy_value =
+          py::reinterpret_borrow<py::object>(py::handle(other_obj));
+      other_tensor = paddle::empty({}, phi::DataType::FLOAT32, place);
+      InitTensorWithNumpyValue(numpy_value, place, &other_tensor);
+    } else {
+      paddle::experimental::Scalar value =
+          CastPyArg2Scalar(other_obj, "__rmod__", 0);
+      if (PyComplex_Check(other_obj)) {
+        eager_gil_scoped_release guard;
+        other_tensor =
+            full_ad_func({1}, value, DataType::COMPLEX64, self_tensor.place());
+      } else {
+        eager_gil_scoped_release guard;
+        other_tensor = full_ad_func(self_tensor.shape(),
+                                    value,
+                                    self_tensor.dtype(),
+                                    self_tensor.place());
+      }
+    }
+    const phi::distributed::ProcessMesh* mesh = nullptr;
+    if (InputsContainDistTensor(&mesh, self_tensor, other_tensor)) {
+      ConvertAllInputsToDistTensor(mesh, self_tensor, other_tensor);
+    }
+  }
+
+  // 3. calculation
+  VLOG(6) << "Calling remainder_ad_func in tensor__rmod__method";
+  {
+    eager_gil_scoped_release guard;
+    ret = remainder_ad_func(other_tensor, self_tensor);
+  }
+  return ToPyObject(ret);
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* tensor__matmul__method(TensorObject* self,
                                         PyObject* args,
                                         PyObject* kwargs) {
@@ -1883,6 +1985,10 @@ PyMethodDef math_op_patch_methods[] = {  // NOLINT
      (PyCFunction)(void (*)())tensor__mod__method,
      METH_VARARGS | METH_KEYWORDS,
      nullptr},
+    {"__rmod__",
+     (PyCFunction)(void (*)())tensor__rmod__method,
+     METH_VARARGS | METH_KEYWORDS,
+     nullptr},
     {"__matmul__",
      (PyCFunction)(void (*)())tensor__matmul__method,
      METH_VARARGS | METH_KEYWORDS,
diff --git a/python/paddle/base/dygraph/math_op_patch.py b/python/paddle/base/dygraph/math_op_patch.py
@@ -231,6 +231,7 @@ def _mT_(var: Tensor) -> Tensor:
         '__rdiv__',
         '__rtruediv__',
         '__mod__',
+        '__rmod__',
         '__matmul__',
         '__gt__',
         '__ge__',
diff --git a/python/paddle/base/layers/math_op_patch.py b/python/paddle/base/layers/math_op_patch.py
@@ -48,6 +48,7 @@
     "__mul__",
     "__rmul__",
     "__mod__",
+    "__rmod__",
     "__div__",
     "__rdiv__",
     "__truediv__",
@@ -78,6 +79,7 @@
     "__rpow__": "A **= B",
     "__floordiv__": "A //B",
     "__mod__": "A % B",
+    "__rmod__": "A %= B",
     "__matmul__": "A @ B",
     "__eq__": "A == B",
     "__ne__": "A != B",
@@ -862,6 +864,10 @@ def to_dense(var):
             '__mod__',
             _binary_creator_('__mod__', 'elementwise_mod', False, None),
         ),
+        (
+            '__rmod__',
+            _binary_creator_('__rmod__', 'elementwise_mod', True, None),
+        ),
         (
             '__matmul__',
             _binary_creator_('__matmul__', "matmul_v2", False, None),
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -44,6 +44,7 @@
     "__mul__",
     "__rmul__",
     "__mod__",
+    "__rmod__",
     "__div__",
     "__rdiv__",
     "__truediv__",
@@ -1095,6 +1096,10 @@ def register_hook(self, hook):
             '__mod__',
             _binary_creator_('__mod__', paddle.tensor.remainder, False, None),
         ),
+        (
+            '__rmod__',
+            _binary_creator_('__rmod__', paddle.tensor.remainder, True, None),
+        ),
         (
             '__matmul__',
             _binary_creator_('__matmul__', paddle.tensor.matmul, False, None),
diff --git a/python/paddle/tensor/tensor.prototype.pyi b/python/paddle/tensor/tensor.prototype.pyi
@@ -176,6 +176,7 @@ class AbstractTensor:
     def __rsub__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rmul__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rtruediv__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
+    def __rmod__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rpow__(self, y: _typing.TensorLike) -> Tensor: ...  # type: ignore
     def __rdiv__(self, y: _typing.TensorLike) -> Tensor: ...
 
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -444,6 +444,15 @@ def test_dygraph_reminder(self):
         np.testing.assert_allclose(actual_out, expect_out, rtol=1e-7, atol=1e-7)
         paddle.enable_static()
 
+    def test_dygraph_rmod(self):
+        paddle.disable_static()
+        self.init_data()
+        # normal case: tenor % nparray
+        expect_out = self.np_a % self.np_b
+        actual_out = self.tensor_b.__rmod__(self.tensor_a)
+        np.testing.assert_allclose(actual_out, expect_out, rtol=1e-7, atol=1e-7)
+        paddle.enable_static()
+
     def test_dygraph_less_than(self):
         paddle.disable_static()
         self.init_data()
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -84,6 +84,7 @@ def test_mod(self):
         res_np_b = x_np % y_np
         res_np_c = paddle.mod(paddle.to_tensor(x_np), paddle.to_tensor(y_np))
         res_np_d = x_np.__mod__(y_np)
+        res_np_e = x_np.__rmod__(y_np)
         paddle.enable_static()
         with paddle.pir_utils.IrGuard():
             main_program, exe, program_guard = new_program()
@@ -97,14 +98,16 @@ def test_mod(self):
                 b = x % y
                 c = x.mod(y)
                 d = x.__mod__(y)
-                (b_np, c_np, d_np) = exe.run(
+                e = x.__rmod__(y)
+                (b_np, c_np, d_np, e_np) = exe.run(
                     main_program,
                     feed={"x": x_np, "y": y_np},
-                    fetch_list=[b, c, d],
+                    fetch_list=[b, c, d, e],
                 )
                 np.testing.assert_allclose(res_np_b, b_np, atol=1e-05)
                 np.testing.assert_allclose(res_np_c, c_np, atol=1e-05)
                 np.testing.assert_allclose(res_np_d, d_np, atol=1e-05)
+                np.testing.assert_allclose(res_np_e, e_np, rtol=1e-05)
 
     def test_matmul(self):
         paddle.disable_static()
