[PHI] fix paddle.Tensor.logit for big tensor (#73046)

DanielSun11 · web-flow · commit 9d2c9e44f730 · 2025-06-05T09:47:15.000+08:00
* fix logit grad

* fix kernel bug  and add unittest

* rm rhol

* fix codestyle error
diff --git a/paddle/phi/kernels/funcs/activation_functor.h b/paddle/phi/kernels/funcs/activation_functor.h
@@ -3375,6 +3375,7 @@ struct CudaLogitGradFunctor : public BaseActivationFunctor<T> {
                ? zero
                : (static_cast<MT>(dout) / (x * (one - x)));
     }
+
     return static_cast<T>(dx);
   }
   static constexpr ActBwdOpFwdDeps FwdDeps() {
diff --git a/test/legacy_test/test_logit_op.py b/test/legacy_test/test_logit_op.py
@@ -86,7 +86,10 @@ def test_check_output(self):
 
     def test_check_grad(self):
         self.check_grad(
-            ['X'], ['Out'], user_defined_grads=[self.x_grad], check_pir=True
+            ['X'],
+            ['Out'],
+            user_defined_grads=[self.x_grad],
+            check_pir=True,
         )
 
 
@@ -99,11 +102,6 @@ def set_attrs(self):
     def test_check_output(self):
         self.check_output(check_pir=True)
 
-    def test_check_grad(self):
-        self.check_grad(
-            ['X'], ['Out'], user_defined_grads=[self.x_grad], check_pir=True
-        )
-
 
 class TestLogitOpFp16(TestLogitOp):
     def set_attrs(self):
@@ -114,11 +112,6 @@ def set_attrs(self):
     def test_check_output(self):
         self.check_output(check_pir=True)
 
-    def test_check_grad(self):
-        self.check_grad(
-            ['X'], ['Out'], user_defined_grads=[self.x_grad], check_pir=True
-        )
-
 
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
@@ -243,6 +236,36 @@ def test_errors(self):
             self.assertRaises(TypeError, paddle.logit, x, dtype='int32')
 
 
+class TestLogitAPI_NAN_Val(unittest.TestCase):
+    def setUp(self):
+        self.init_input_output()
+        self.place = [paddle.CPUPlace()]
+        if paddle.base.core.is_compiled_with_cuda():
+            self.place.append(paddle.CUDAPlace(0))
+
+    def init_input_output(self):
+        self.x = [-0.1, 1.1, 2]
+        self.expect_out = [np.nan, np.nan, np.nan]
+        self.expect_x_grad = [np.nan, np.nan, np.nan]
+
+    def test_nan_val(self):
+        def _test_nan_val_with_place(place):
+            with paddle.base.dygraph.guard():
+                x = paddle.to_tensor(self.x, stop_gradient=False, place=place)
+                y = paddle.logit(x)
+                loss = y.sum()
+                loss.backward()
+                np.testing.assert_allclose(
+                    y.numpy(), self.expect_out, rtol=1e-05
+                )
+                np.testing.assert_allclose(
+                    x.grad.numpy(), self.expect_x_grad, rtol=1e-05
+                )
+
+        for place in self.place:
+            _test_nan_val_with_place(place)
+
+
 class TestLogitAPICase1(unittest.TestCase):
     def init_data(self):
         self.x_shape = [120]

Original file line number	Diff line number	Diff line change
`@@ -3375,6 +3375,7 @@ struct CudaLogitGradFunctor : public BaseActivationFunctor<T> {`
`3375`	`3375`	`? zero`
`3376`	`3376`	`: (static_cast<MT>(dout) / (x * (one - x)));`
`3377`	`3377`	`}`
	`3378`	`+`
`3378`	`3379`	`return static_cast<T>(dx);`
`3379`	`3380`	`}`
`3380`	`3381`	`static constexpr ActBwdOpFwdDeps FwdDeps() {`