[Feat] Support custom __radd__ and __rmul__ (#73119)

BeingGod · web-flow · commit d5bd8000456c · 2025-06-09T21:38:48.000+08:00
* [Feat] Support custom __radd__ and __rmul__

* resolve reviewer issues
diff --git a/paddle/fluid/pybind/eager_math_op_patch.cc b/paddle/fluid/pybind/eager_math_op_patch.cc
@@ -330,7 +330,12 @@ static PyObject* tensor__add__method(TensorObject* self,
   }
 
   return ToPyObject(ret);
+
+#if defined(PADDLE_WITH_XPU_BKCL)
   EAGER_CATCH_AND_THROW_RETURN_NULL
+#else
+  EAGER_CATCH_AND_THROW_RETURN_NOT_IMPLEMENTED
+#endif
 }
 
 static PyObject* tensor__sub__method(TensorObject* self,
@@ -625,7 +630,12 @@ static PyObject* tensor__mul__method(TensorObject* self,
   }
 
   return ToPyObject(ret);
+
+#if defined(PADDLE_WITH_XPU_BKCL)
   EAGER_CATCH_AND_THROW_RETURN_NULL
+#else
+  EAGER_CATCH_AND_THROW_RETURN_NOT_IMPLEMENTED
+#endif
 }
 
 static PyObject* tensor__div__method(TensorObject* self,
diff --git a/paddle/fluid/pybind/exception.h b/paddle/fluid/pybind/exception.h
@@ -33,6 +33,13 @@ limitations under the License. */
     return -1;                                        \
   }
 
+#define EAGER_CATCH_AND_THROW_RETURN_NOT_IMPLEMENTED \
+  }                                                  \
+  catch (...) {                                      \
+    Py_INCREF(Py_NotImplemented);                    \
+    return Py_NotImplemented;                        \
+  }
+
 namespace paddle {
 namespace pybind {
 
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -503,6 +503,20 @@ def test_builtin_type_conversion(self):
             complex(a)
 
 
+class CustomTensor:
+
+    def __init__(self, value):
+        assert isinstance(value, paddle.Tensor)
+
+        self._value = value
+
+    def __radd__(self, other):
+        return self._value + other
+
+    def __rmul__(self, other):
+        return self._value * other
+
+
 class TestDygraphMathOpPatches(unittest.TestCase):
     def init_data(self):
         self.np_a = np.random.random((2, 3, 4)).astype(np.float32)
@@ -636,6 +650,34 @@ def test_dygraph_rmatmul(self):
         np.testing.assert_allclose(a @ b, c.numpy(), rtol=1e-5, atol=1e-5)
         paddle.enable_static()
 
+    @unittest.skipIf(
+        paddle.device.is_compiled_with_xpu(),
+        reason="XPU not support custom tensor.",
+    )
+    def test_dygraph_custom_add(self):
+        paddle.disable_static()
+        a_np = np.random.random((2, 3)).astype(np.float32) * 100
+        b_np = np.random.random((2, 3)).astype(np.float32) * 100
+        a = paddle.to_tensor(a_np)
+        b = paddle.to_tensor(b_np)
+        c = a + CustomTensor(b)
+        np.testing.assert_allclose(a + b, c, atol=0)
+        paddle.enable_static()
+
+    @unittest.skipIf(
+        paddle.device.is_compiled_with_xpu(),
+        reason="XPU not support custom tensor.",
+    )
+    def test_dygraph_custom_mul(self):
+        paddle.disable_static()
+        a_np = np.random.random((2, 3)).astype(np.float32) * 100
+        b_np = np.random.random((2, 3)).astype(np.float32) * 100
+        a = paddle.to_tensor(a_np)
+        b = paddle.to_tensor(b_np)
+        c = a * CustomTensor(b)
+        np.testing.assert_allclose(a * b, c, atol=0)
+        paddle.enable_static()
+
 
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -330,7 +330,12 @@ static PyObject* tensor__add__method(TensorObject* self,`
`330`	`330`	`}`
`331`	`331`
`332`	`332`	`return ToPyObject(ret);`
	`333`	`+`
	`334`	`+#if defined(PADDLE_WITH_XPU_BKCL)`
`333`	`335`	`EAGER_CATCH_AND_THROW_RETURN_NULL`
	`336`	`+#else`
	`337`	`+ EAGER_CATCH_AND_THROW_RETURN_NOT_IMPLEMENTED`
	`338`	`+#endif`
`334`	`339`	`}`
`335`	`340`
`336`	`341`	`static PyObject* tensor__sub__method(TensorObject* self,`
`@@ -625,7 +630,12 @@ static PyObject* tensor__mul__method(TensorObject* self,`
`625`	`630`	`}`
`626`	`631`
`627`	`632`	`return ToPyObject(ret);`
	`633`	`+`
	`634`	`+#if defined(PADDLE_WITH_XPU_BKCL)`
`628`	`635`	`EAGER_CATCH_AND_THROW_RETURN_NULL`
	`636`	`+#else`
	`637`	`+ EAGER_CATCH_AND_THROW_RETURN_NOT_IMPLEMENTED`
	`638`	`+#endif`
`629`	`639`	`}`
`630`	`640`
`631`	`641`	`static PyObject* tensor__div__method(TensorObject* self,`