[PIR Unittest] fix pir ut (test_while_op,test_norm_nn_grad) (#66785)

* [Fix PIR Unittest] fix pir ut (test_while_op,test_norm_nn_grad)

* update

* fix feed dtype

* add dtype infer for ShuffleBatchInferMeta

* update place

* fix ci coverage

paddle/fluid/pir/dialect/op_generator/ops_api_gen.py

@@ -212,7 +212,6 @@
     'seed',
     'shadow_feed',
     'shadow_feed_tensors',
-    'shuffle_batch',
     'sparse_momentum',
     'tdm_sampler',
     'soft_relu',

paddle/phi/infermeta/binary.cc

@@ -3759,8 +3759,10 @@ void ShuffleBatchInferMeta(const MetaTensor& x,
 ) {
   out->share_dims(x);
   out->share_lod(x);
+  out->set_dtype(x.dtype());
   seed_out->share_dims(seed);
   seed_out->share_lod(seed);
+  seed_out->set_dtype(seed.dtype());
   shuffle_idx->set_dims(phi::make_ddim({-1}));
 }
 

python/paddle/incubate/layers/nn.py

@@ -530,11 +530,19 @@ def shuffle_batch(x: Tensor, seed: int | Tensor | None = None) -> Tensor:
     op_attrs = {}
     if isinstance(seed, int):
         op_attrs["startup_seed"] = seed
-        seed = helper.create_variable(
-            name=unique_name.generate("shuffle_batch_seed"),
-            dtype="int64",
-            persistable=False,
-        )
+        if paddle.framework.in_pir_mode():
+            seed = paddle.full([0], 0, "int64")
+            out, _, _ = _C_ops.shuffle_batch(x, seed, op_attrs["startup_seed"])
+            return out
+        else:
+            seed = helper.create_variable(
+                name=unique_name.generate("shuffle_batch_seed"),
+                dtype="int64",
+                persistable=False,
+            )
+    if paddle.framework.in_pir_mode():
+        out, _, _ = _C_ops.shuffle_batch(x, seed, 0)
+        return out
     helper.append_op(
         type='shuffle_batch',
         inputs={'X': x, 'Seed': seed},

test/legacy_test/test_norm_nn_grad.py

@@ -71,7 +71,8 @@ def test_grad(self):
         if core.is_compiled_with_cuda():
             places.append(base.CUDAPlace(0))
         for p in places:
-            self.func(p)
+            with paddle.pir_utils.OldIrGuard():
+                self.func(p)
             self.func_pir(p)
 
 
@@ -281,7 +282,8 @@ def test_grad(self):
         if core.is_compiled_with_cuda():
             places.append(base.CUDAPlace(0))
         for p in places:
-            self.func(p)
+            with paddle.pir_utils.OldIrGuard():
+                self.func(p)
             self.func_pir(p)
 
 

test/legacy_test/test_shuffle_batch_op.py

@@ -19,7 +19,9 @@
 import numpy as np
 from op_test import OpTest
 
+import paddle
 from paddle import base
+from paddle.pir_utils import test_with_pir_api
 
 
 class TestShuffleBatchOpBase(OpTest):
@@ -38,6 +40,8 @@ def _get_places(self):
 
     def setUp(self):
         self.op_type = 'shuffle_batch'
+        self.python_api = paddle.incubate.layers.shuffle_batch
+        self.python_out_sig = ["Out"]
         self.dtype = np.float64
         self.shape = self.get_shape()
         x = self.gen_random_array(self.shape)
@@ -53,7 +57,7 @@ def setUp(self):
         self.attrs = {'startup_seed': 1}
 
     def test_check_output(self):
-        self.check_output_customized(self.verify_output)
+        self.check_output_customized(self.verify_output, check_pir=True)
 
     def verify_output(self, outs):
         x = np.copy(self.inputs['X'])
@@ -76,13 +80,65 @@ def sort_array(self, array):
         return np.reshape(np.array(arr_list), shape)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_dygraph=False)
+        self.check_grad(['X'], 'Out', check_dygraph=False, check_pir=True)
 
 
 class TestShuffleBatchOp2(TestShuffleBatchOpBase):
     def get_shape(self):
         return (4, 30)
 
 
+class TestShuffleBatchAPI(unittest.TestCase):
+    def setUp(self):
+        self.places = [paddle.CPUPlace()]
+        if not os.name == 'nt' and paddle.is_compiled_with_cuda():
+            self.places.append(paddle.CUDAPlace(0))
+        paddle.enable_static()
+
+    def tearDown(self):
+        paddle.disable_static()
+
+    @test_with_pir_api
+    def test_seed_without_tensor(self):
+        def api_run(seed, place=paddle.CPUPlace()):
+            main_prog, startup_prog = (
+                paddle.static.Program(),
+                paddle.static.Program(),
+            )
+            with paddle.static.program_guard(main_prog, startup_prog):
+                x = paddle.static.data(name='x', shape=[-1, 4], dtype='float32')
+                out = paddle.incubate.layers.shuffle_batch(x, seed=seed)
+            exe = paddle.static.Executor(place=place)
+            feed = {'x': np.random.random((10, 4)).astype('float32')}
+            exe.run(startup_prog)
+            _ = exe.run(main_prog, feed=feed, fetch_list=[out])
+
+        for place in self.places:
+            api_run(None, place=place)
+            api_run(1, place=place)
+
+    @test_with_pir_api
+    def test_seed_with_tensor(self):
+        def api_run(place=paddle.CPUPlace()):
+            main_prog, startup_prog = (
+                paddle.static.Program(),
+                paddle.static.Program(),
+            )
+            with paddle.static.program_guard(main_prog, startup_prog):
+                x = paddle.static.data(name='x', shape=[-1, 4], dtype='float32')
+                seed = paddle.static.data(name='seed', shape=[1], dtype='int64')
+                out = paddle.incubate.layers.shuffle_batch(x, seed=seed)
+            exe = paddle.static.Executor(place=place)
+            feed = {
+                'x': np.random.random((10, 4)).astype('float32'),
+                'seed': np.array([1]).astype('int64'),
+            }
+            exe.run(startup_prog)
+            _ = exe.run(main_prog, feed=feed, fetch_list=[out])
+
+        for place in self.places:
+            api_run(place=place)
+
+
 if __name__ == '__main__':
     unittest.main()

test/legacy_test/test_while_op.py

@@ -159,6 +159,7 @@ def test_bad_x():
 
 
 class TestIgnoreVarNameInWhile(unittest.TestCase):
+    @test_with_pir_api
     def test_ignore_var(self):
         def cond(i, ten, temp, y):
             return i < ten
@@ -189,9 +190,11 @@ def body_func(i, ten, batch_info, origin_seq):
         exe = base.Executor(base.CPUPlace())
         exe.run(base.default_startup_program())
 
-        input_x = numpy.array([[1, 2, 3, 4], [4, 5, 6, 7], [7, 8, 9, 10]])
+        input_x = numpy.array(
+            [[1.0, 2.0, 3.0, 4.0], [4.0, 5.0, 6.0, 7.0], [7.0, 8.0, 9.0, 10.0]]
+        ).astype('float32')
         input_x = input_x.reshape(3, 1, 4)
-        input_y = numpy.array([[10], [12], [33]])
+        input_y = numpy.array([[10.0], [12.0], [33.0]]).astype('float32')
         input_y = input_y.reshape(3, 1, 1)
 
         (res,) = exe.run(