[PIR] Add missing assign for divide with scalar (#62252)

SigureMo · web-flow · commit 754079f9df70 · 2024-03-01T09:49:35.000+08:00
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -338,7 +338,7 @@ def __impl__(self, other_var):
                     python_api == paddle.divide
                     and self.dtype in _supported_int_dtype_
                 ):
-                    paddle.cast(self, DataType.FLOAT32)
+                    self = paddle.cast(self, DataType.FLOAT32)
                 # here use `scale` replace `elementwise` to get better performance
                 # but only +, -, *, / can use this method
                 if scalar_method is not None:
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -643,6 +643,32 @@ def test_math_exists(self):
             self.assertTrue(inspect.ismethod(a.asinh_))
             self.assertTrue(inspect.ismethod(a.diag))
 
+    def test_binary_op_with_scalar(self):
+        with paddle.pir_utils.IrGuard():
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x_np = np.array(10, dtype=np.int32)
+                x = paddle.static.data(name='x', shape=[], dtype="int32")
+                y1 = x / 2
+                y2 = x / 5.0
+                y3 = x // 2
+                y4 = x * 8.0
+                self.assertEqual(y1.dtype, paddle.pir.core.DataType.FLOAT32)
+                self.assertEqual(y2.dtype, paddle.pir.core.DataType.FLOAT32)
+                self.assertEqual(y3.dtype, paddle.pir.core.DataType.INT32)
+                self.assertEqual(y4.dtype, paddle.pir.core.DataType.FLOAT32)
+                (y1_out, y2_out, y3_out, y4_out) = exe.run(
+                    main_program,
+                    feed={
+                        "x": x_np,
+                    },
+                    fetch_list=[y1, y2, y3, y4],
+                )
+                np.testing.assert_allclose(x_np / 2, y1_out, rtol=1e-05)
+                np.testing.assert_allclose(x_np / 5.0, y2_out, rtol=1e-05)
+                np.testing.assert_allclose(x_np // 2, y3_out, atol=1e-05)
+                np.testing.assert_allclose(x_np * 8.0, y4_out, rtol=1e-05)
+
 
 if __name__ == '__main__':
     unittest.main()
