Fix argsort cpu kernel when with input of NaN

paddle/phi/kernels/cpu/argsort_kernel.cc

@@ -51,9 +51,13 @@ static void FullSort(Type input_height,
               col_vec.end(),
               [&](const std::pair<T, Type>& l, const std::pair<T, Type>& r) {
                 if (descending)
-                  return l.first > r.first;
+                  return (std::isnan(static_cast<double>(l.first)) &&
+                          !std::isnan(static_cast<double>(r.first))) ||
+                         (l.first > r.first);
                 else
-                  return l.first < r.first;
+                  return (!std::isnan(static_cast<double>(l.first)) &&
+                          std::isnan(static_cast<double>(r.first))) ||
+                         (l.first < r.first);
               });
 
     for (Type j = 0; j < input_width; ++j) {

python/paddle/fluid/tests/unittests/test_argsort_op.py

@@ -442,5 +442,28 @@ def init(self):
         self.axis = 1
 
 
+class TestArgsortWithInputNaN(unittest.TestCase):
+    def init(self):
+        self.axis = 0
+
+    def setUp(self):
+        self.init()
+        self.input_data = np.array([1.0, np.nan, 3.0, 2.0])
+        if core.is_compiled_with_cuda():
+            self.place = core.CUDAPlace(0)
+        else:
+            self.place = core.CPUPlace()
+
+    def test_api(self):
+        paddle.disable_static(self.place)
+        var_x = paddle.to_tensor(self.input_data)
+        out = paddle.argsort(var_x, axis=self.axis)
+        self.assertEqual((out.numpy() == np.array([0, 3, 2, 1])).all(), True)
+
+        out = paddle.argsort(var_x, axis=self.axis, descending=True)
+        self.assertEqual((out.numpy() == np.array([1, 2, 3, 0])).all(), True)
+        paddle.enable_static()
+
+
 if __name__ == "__main__":
     unittest.main()