From 85df6d73b9ea42967649ca7426c526c2f548c9a1 Mon Sep 17 00:00:00 2001
From: yeliang2258 <30516196+yeliang2258@users.noreply.github.com>
Date: Mon, 8 Aug 2022 21:02:25 +0800
Subject: [PATCH] Add inf and nan support in equal OP (#44667)

* add inf and nan support in equal

* add header

* fix nan and update test

* update test

* update test

* update test

* update code

* update compare test

* update func

* update

* update

* fix

* update
---
 paddle/phi/kernels/funcs/compare_functors.h   |  5 +
 .../fluid/tests/unittests/test_compare_op.py  | 96 +++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/paddle/phi/kernels/funcs/compare_functors.h b/paddle/phi/kernels/funcs/compare_functors.h
index 569fed7b7f..e16083506b 100644
--- a/paddle/phi/kernels/funcs/compare_functors.h
+++ b/paddle/phi/kernels/funcs/compare_functors.h
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+#include <math.h>
 
 namespace phi {
 namespace funcs {
@@ -35,6 +36,10 @@ template <typename InT, typename OutT = bool>
 struct EqualFunctor {
   HOSTDEVICE OutT operator()(const InT a, const InT b) const {
     if (std::is_floating_point<InT>::value) {
+      if (isinf(static_cast<float>(a)) || isinf(static_cast<float>(b)))
+        return static_cast<OutT>(a == b);
+      if (isnan(static_cast<float>(a)) || isnan(static_cast<float>(b)))
+        return static_cast<OutT>(false);
       return static_cast<OutT>(fabs(static_cast<double>(a - b)) < 1e-8);
     } else {
       return static_cast<OutT>(a == b);
diff --git a/python/paddle/fluid/tests/unittests/test_compare_op.py b/python/paddle/fluid/tests/unittests/test_compare_op.py
index f1bacfbb6f..731eedfca6 100755
--- a/python/paddle/fluid/tests/unittests/test_compare_op.py
+++ b/python/paddle/fluid/tests/unittests/test_compare_op.py
@@ -150,6 +150,102 @@ def create_paddle_case(op_type, callback):
                 self.assertEqual((out.numpy() == self.real_result).all(), True)
                 paddle.enable_static()
 
+        def test_dynamic_api_inf_1(self):
+            if self.op_type == "equal":
+                paddle.disable_static()
+                x1 = np.array([1, float('inf'), float('inf')]).astype(np.int64)
+                x = paddle.to_tensor(x1)
+                y1 = np.array([1, float('-inf'), float('inf')]).astype(np.int64)
+                y = paddle.to_tensor(y1)
+                op = eval("paddle.%s" % (self.op_type))
+                out = op(x, y)
+                self.real_result = (x1 == y1).astype(np.int64)
+                self.assertEqual(
+                    (out.numpy().astype(np.int64) == self.real_result).all(),
+                    True)
+                paddle.enable_static()
+
+        def test_dynamic_api_inf_2(self):
+            if self.op_type == "equal":
+                paddle.disable_static()
+                x1 = np.array([1, float('inf'),
+                               float('inf')]).astype(np.float32)
+                x = paddle.to_tensor(x1)
+                y1 = np.array([1, float('-inf'),
+                               float('inf')]).astype(np.float32)
+                y = paddle.to_tensor(y1)
+                op = eval("paddle.%s" % (self.op_type))
+                out = op(x, y)
+                self.real_result = (x1 == y1).astype(np.int64)
+                self.assertEqual(
+                    (out.numpy().astype(np.int64) == self.real_result).all(),
+                    True)
+                paddle.enable_static()
+
+        def test_dynamic_api_inf_3(self):
+            if self.op_type == "equal":
+                paddle.disable_static()
+                x1 = np.array([1, float('inf'),
+                               float('-inf')]).astype(np.float32)
+                x = paddle.to_tensor(x1)
+                y1 = np.array([1, 2, 3]).astype(np.float32)
+                y = paddle.to_tensor(y1)
+                op = eval("paddle.%s" % (self.op_type))
+                out = op(x, y)
+                self.real_result = (x1 == y1).astype(np.int64)
+                self.assertEqual(
+                    (out.numpy().astype(np.int64) == self.real_result).all(),
+                    True)
+                paddle.enable_static()
+
+        def test_dynamic_api_nan_1(self):
+            if self.op_type == "equal":
+                paddle.disable_static()
+                x1 = np.array([1, float('nan'), float('nan')]).astype(np.int64)
+                x = paddle.to_tensor(x1)
+                y1 = np.array([1, float('-nan'), float('nan')]).astype(np.int64)
+                y = paddle.to_tensor(y1)
+                op = eval("paddle.%s" % (self.op_type))
+                out = op(x, y)
+                self.real_result = (x1 == y1).astype(np.int64)
+                self.assertEqual(
+                    (out.numpy().astype(np.int64) == self.real_result).all(),
+                    True)
+                paddle.enable_static()
+
+        def test_dynamic_api_nan_2(self):
+            if self.op_type == "equal":
+                paddle.disable_static()
+                x1 = np.array([1, float('nan'),
+                               float('nan')]).astype(np.float32)
+                x = paddle.to_tensor(x1)
+                y1 = np.array([1, float('-nan'),
+                               float('nan')]).astype(np.float32)
+                y = paddle.to_tensor(y1)
+                op = eval("paddle.%s" % (self.op_type))
+                out = op(x, y)
+                self.real_result = (x1 == y1).astype(np.int64)
+                self.assertEqual(
+                    (out.numpy().astype(np.int64) == self.real_result).all(),
+                    True)
+                paddle.enable_static()
+
+        def test_dynamic_api_nan_3(self):
+            if self.op_type == "equal":
+                paddle.disable_static()
+                x1 = np.array([1, float('-nan'),
+                               float('nan')]).astype(np.float32)
+                x = paddle.to_tensor(x1)
+                y1 = np.array([1, 2, 1]).astype(np.float32)
+                y = paddle.to_tensor(y1)
+                op = eval("paddle.%s" % (self.op_type))
+                out = op(x, y)
+                self.real_result = (x1 == y1).astype(np.int64)
+                self.assertEqual(
+                    (out.numpy().astype(np.int64) == self.real_result).all(),
+                    True)
+                paddle.enable_static()
+
         def test_not_equal(self):
             if self.op_type == "not_equal":
                 paddle.disable_static()
-- 
GitLab