Support builtin __abs__ method for Tensor/Variable/Value (#68511)

HydrogenSulfate · web-flow · commit 21ec722e28ec · 2024-10-08T10:40:31.000+08:00
diff --git a/python/paddle/base/dygraph/math_op_patch.py b/python/paddle/base/dygraph/math_op_patch.py
@@ -107,6 +107,9 @@ def _scalar_elementwise_op_(
     def _neg_(var: Tensor) -> Tensor:
         return _scalar_elementwise_op_(var, -1.0, 0.0)
 
+    def _abs_(var: Tensor) -> Tensor:
+        return var.abs()
+
     def _float_(var: Tensor) -> float:
         numel = np.prod(var.shape)
         assert (
@@ -188,6 +191,7 @@ def _T_(var: Tensor) -> Tensor:
 
     eager_methods = [
         ('__neg__', _neg_),
+        ('__abs__', _abs_),
         ('__float__', _float_),
         ('__long__', _long_),
         ('__int__', _int_),
diff --git a/python/paddle/base/layers/math_op_patch.py b/python/paddle/base/layers/math_op_patch.py
@@ -482,6 +482,9 @@ def _scalar_op_(var, scale, bias):
     def _neg_(var):
         return _scalar_op_(var, -1.0, 0.0)
 
+    def _abs_(var):
+        return paddle.abs(var)
+
     @property
     def _ndim(self):
         """
@@ -778,6 +781,7 @@ def to_dense(var):
     variable_methods = [
         #   b=-a
         ('__neg__', _neg_),
+        ('__abs__', _abs_),
         ('astype', astype),
         ('cpu', cpu),
         ('cuda', cuda),
diff --git a/python/paddle/optimizer/lbfgs.py b/python/paddle/optimizer/lbfgs.py
@@ -209,7 +209,7 @@ def _strong_wolfe(
             bracket_gtd = [gtd_prev, gtd_new]
             break
 
-        if paddle.abs(gtd_new) <= -c2 * gtd:
+        if abs(gtd_new) <= -c2 * gtd:
             bracket = [alpha]
             bracket_f = [loss_new]
             bracket_g = [grad_new]
@@ -262,10 +262,7 @@ def _strong_wolfe(
     low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
     while not done and ls_iter < max_ls:
         # line-search bracket is so small
-        bracket_ls = bracket[1] - bracket[0]
-        if not isinstance(bracket_ls, paddle.Tensor):
-            bracket_ls = paddle.to_tensor(bracket_ls, dtype=gtd_new.dtype)
-        if paddle.abs(bracket_ls) * d_norm < tolerance_change:
+        if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
             break
 
         # compute new trial value
@@ -291,9 +288,7 @@ def _strong_wolfe(
             # interpolation close to boundary
             if insuf_progress or alpha >= max(bracket) or alpha <= min(bracket):
                 # evaluate at 0.1 away from boundary
-                if paddle.abs(alpha - max(bracket)) < paddle.abs(
-                    alpha - min(bracket)
-                ):
+                if abs(alpha - max(bracket)) < abs(alpha - min(bracket)):
                     alpha = max(bracket) - eps
                 else:
                     alpha = min(bracket) + eps
@@ -321,7 +316,7 @@ def _strong_wolfe(
                 (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
             )
         else:
-            if paddle.abs(gtd_new) <= -c2 * gtd:
+            if abs(gtd_new) <= -c2 * gtd:
                 # Wolfe conditions satisfied
                 done = True
             elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -334,6 +334,9 @@ def _scalar_div_(var, value):
     def _scalar_neg_(var):
         return paddle.scale(var, -1.0, 0.0)
 
+    def _scalar_abs_(var):
+        return paddle.abs(var)
+
     def _binary_creator_(
         method_name,
         python_api,
@@ -1063,6 +1066,7 @@ def register_hook(self, hook):
             _binary_creator_('__matmul__', paddle.tensor.matmul, False, None),
         ),
         ('__neg__', _scalar_neg_),
+        ('__abs__', _scalar_abs_),
         # For compare operators
         (
             '__eq__',
diff --git a/test/legacy_test/test_math_op_patch.py b/test/legacy_test/test_math_op_patch.py
@@ -245,6 +245,40 @@ def test_neg(self):
         )
         np.testing.assert_allclose(-a_np, b_np, rtol=1e-05)
 
+    @prog_scope()
+    def test_abs(self):
+        # test for real number
+        a = paddle.static.data(name="a", shape=[-1, 10, 1], dtype='float32')
+        if not paddle.framework.use_pir_api():
+            a.desc.set_need_check_feed(False)
+        b = abs(a)  # call __abs__
+        place = base.CPUPlace()
+        exe = base.Executor(place)
+        a_np = np.random.uniform(-1, 1, size=[10, 1]).astype('float32')
+
+        (b_np,) = exe.run(
+            base.default_main_program(), feed={"a": a_np}, fetch_list=[b]
+        )
+        np.testing.assert_allclose(np.abs(a_np), b_np, rtol=1e-05)
+
+    @prog_scope()
+    def test_abs_complex(self):
+        # test for complex number
+        a = paddle.static.data(name="a", shape=[-1, 10, 1], dtype='complex64')
+        if not paddle.framework.use_pir_api():
+            a.desc.set_need_check_feed(False)
+        b = abs(a)  # call __abs__
+        place = base.CPUPlace()
+        exe = base.Executor(place)
+        a_np = np.random.uniform(-1, 1, size=[10, 1]).astype(
+            'float32'
+        ) + 1j * np.random.uniform(-1, 1, size=[10, 1]).astype('float32')
+
+        (b_np,) = exe.run(
+            base.default_main_program(), feed={"a": a_np}, fetch_list=[b]
+        )
+        np.testing.assert_allclose(np.abs(a_np), b_np, rtol=1e-05)
+
     @prog_scope()
     def test_astype(self):
         a = paddle.static.data(name="a", shape=[-1, 10, 1])
diff --git a/test/legacy_test/test_math_op_patch_pir.py b/test/legacy_test/test_math_op_patch_pir.py
@@ -554,6 +554,48 @@ def test_neg(self):
                 np.testing.assert_array_equal(res, a_np)
                 np.testing.assert_array_equal(res, b_np)
 
+    def test_abs(self):
+        # test for real number
+        x_np = np.random.uniform(-1, 1, [10, 1024]).astype(np.float32)
+        res = abs(x_np)
+        with paddle.pir_utils.IrGuard():
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.static.data(
+                    name='x', shape=[10, 1024], dtype="float32"
+                )
+                a = abs(x)
+                b = x.__abs__()
+                (a_np, b_np) = exe.run(
+                    main_program,
+                    feed={"x": x_np},
+                    fetch_list=[a, b],
+                )
+                np.testing.assert_array_equal(res, a_np)
+                np.testing.assert_array_equal(res, b_np)
+
+    def test_abs_complex(self):
+        # test for complex number
+        x_np = np.random.uniform(-1, 1, [10, 1024]).astype(
+            np.float32
+        ) + 1j * np.random.uniform(-1, 1, [10, 1024]).astype(np.float32)
+        res = abs(x_np)
+        with paddle.pir_utils.IrGuard():
+            main_program, exe, program_guard = new_program()
+            with program_guard:
+                x = paddle.static.data(
+                    name='x', shape=[10, 1024], dtype="complex64"
+                )
+                a = abs(x)
+                b = x.__abs__()
+                (a_np, b_np) = exe.run(
+                    main_program,
+                    feed={"x": x_np},
+                    fetch_list=[a, b],
+                )
+                np.testing.assert_allclose(res, a_np, rtol=2e-7, atol=0.0)
+                np.testing.assert_allclose(res, b_np, rtol=2e-7, atol=0.0)
+
     def test_builtin_type_conversion(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
@@ -479,6 +479,26 @@ def test_neg(self):
             res = -a
             np.testing.assert_array_equal(res.numpy(), -a_np)
 
+    def test_abs(self):
+        # test for real number
+        a_np = np.random.uniform(-1, 1, self.shape).astype(self.dtype)
+        with base.dygraph.guard():
+            a = paddle.to_tensor(a_np)
+            res = abs(a)
+            np.testing.assert_array_equal(res.numpy(), np.abs(a_np))
+
+    def test_abs_complex(self):
+        # test for complex number
+        a_np = np.random.uniform(-1, 1, self.shape).astype(
+            self.dtype
+        ) + 1j * np.random.uniform(-1, 1, self.shape).astype(self.dtype)
+        with base.dygraph.guard():
+            a = paddle.to_tensor(a_np)
+            res = abs(a)
+            np.testing.assert_allclose(
+                res.numpy(), np.abs(a_np), rtol=2e-7, atol=0.0
+            )
+
     def test_float_int_long(self):
         with base.dygraph.guard():
             a = paddle.to_tensor(np.array([100.1]))
