fix nonzero (#72003)

aquagull · web-flow · commit 4368b7f8e268 · 2025-04-03T18:31:31.000+08:00
* fix

* fix test

* fix
diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py
@@ -523,10 +523,6 @@ def nonzero(x: Tensor, as_tuple=False):
              [3]])
 
     """
-    list_out = []
-    shape = x.shape
-    rank = len(shape)
-
     if in_dynamic_or_pir_mode():
         outs = _C_ops.nonzero(x)
     else:
@@ -558,13 +554,9 @@ def nonzero(x: Tensor, as_tuple=False):
 
     if not as_tuple:
         return outs
-    elif rank == 1:
-        return (outs,)
     else:
-        for i in range(rank):
-            list_out.append(
-                paddle.slice(outs, axes=[1], starts=[i], ends=[i + 1])
-            )
+        rank = x.ndim
+        list_out = [outs[:, i] for i in range(rank)]
         return tuple(list_out)
 
 
diff --git a/test/legacy_test/test_nonzero_api.py b/test/legacy_test/test_nonzero_api.py
@@ -38,13 +38,13 @@ def test_nonzero_api_as_tuple(self):
             y = paddle.nonzero(x, as_tuple=True)
             self.assertEqual(type(y), tuple)
             self.assertEqual(len(y), 2)
-            z = paddle.concat(list(y), axis=1)
+            z = paddle.concat(list(y), axis=0)
             exe = base.Executor(base.CPUPlace())
 
             (res,) = exe.run(
                 feed={'x': data}, fetch_list=[z], return_numpy=False
             )
-        expect_out = np.array([[0, 0], [1, 1]])
+        expect_out = np.array([0, 1, 0, 1])
         np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
 
         data = np.array([1, 1, 0], dtype="float32")
@@ -55,12 +55,12 @@ def test_nonzero_api_as_tuple(self):
             y = paddle.nonzero(x, as_tuple=True)
             self.assertEqual(type(y), tuple)
             self.assertEqual(len(y), 1)
-            z = paddle.concat(list(y), axis=1)
+            z = paddle.concat(list(y), axis=0)
             exe = base.Executor(base.CPUPlace())
             (res,) = exe.run(
                 feed={'x': data}, fetch_list=[z], return_numpy=False
             )
-        expect_out = np.array([[0], [1]])
+        expect_out = np.array([0, 1])
         np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
 
     def test_nonzero_api(self):
diff --git a/test/legacy_test/test_where_op.py b/test/legacy_test/test_where_op.py
@@ -788,12 +788,12 @@ def test_where_condition(self):
             y = paddle.where(x)
             self.assertEqual(type(y), tuple)
             self.assertEqual(len(y), 2)
-            z = paddle.concat(list(y), axis=1)
+            z = paddle.concat(list(y), axis=0)
             exe = base.Executor(base.CPUPlace())
             (res,) = exe.run(
                 feed={'x': data}, fetch_list=[z], return_numpy=False
             )
-        expect_out = np.array([[0, 0], [1, 1]])
+        expect_out = np.array([0, 1, 0, 1])
         np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
         data = np.array([True, True, False])
         with program_guard(Program(), Program()):
@@ -803,12 +803,12 @@ def test_where_condition(self):
             y = paddle.where(x)
             self.assertEqual(type(y), tuple)
             self.assertEqual(len(y), 1)
-            z = paddle.concat(list(y), axis=1)
+            z = paddle.concat(list(y), axis=0)
             exe = base.Executor(base.CPUPlace())
             (res,) = exe.run(
                 feed={'x': data}, fetch_list=[z], return_numpy=False
             )
-        expect_out = np.array([[0], [1]])
+        expect_out = np.array([0, 1])
         np.testing.assert_allclose(expect_out, np.array(res), rtol=1e-05)
 
 
diff --git a/test/xpu/test_where_index_xpu.py b/test/xpu/test_where_index_xpu.py
@@ -112,7 +112,7 @@ def test_errors(self):
         def test_type():
             paddle.nonzero([10])
 
-        self.assertRaises(AttributeError, test_type)
+        self.assertRaises(TypeError, test_type)
 
 
 class TestWhereSimulatorMode(unittest.TestCase):
