[PIR]Migrate gather into pir (#58498)

changeyoung98 · 0x45f · web-flow · commit bc76ee063506 · 2023-10-31T10:44:38.000+08:00
* [PIR]Migrate gather into pir

* Open UT

* fix bug in TestGatherOp_ZeroDim

---------

Co-authored-by: 0x45f &lt;wangzhen45@baidu.com&gt;
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -1545,6 +1545,8 @@ void GatherInferMeta(const MetaTensor& x,
     if (input_dim.size() == 1) {
       // the index is a 0D tensor and the x is a 1D tensor
       out->set_dims(phi::DDim(phi::Dim<0>()));
+      out->set_dtype(x.dtype());
+      out->share_lod(x);
     } else {
       if (axis.FromTensor() || axis_v == 0) {
         // decrease the output dimension
diff --git a/test/legacy_test/test_gather_op.py b/test/legacy_test/test_gather_op.py
@@ -42,10 +42,10 @@ def setUp(self):
         self.if_enable_cinn()
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True)
+        self.check_grad(['X'], 'Out', check_prim=True, check_pir=True)
 
     def config(self):
         """
@@ -115,11 +115,11 @@ def if_enable_cinn(self):
         self.enable_cinn = False
 
     def test_check_output(self):
-        self.check_output_with_place(place=paddle.CUDAPlace(0))
+        self.check_output_with_place(place=paddle.CUDAPlace(0), check_pir=True)
 
     def test_check_grad(self):
         self.check_grad_with_place(
-            paddle.CUDAPlace(0), ['X'], 'Out', check_prim=True
+            paddle.CUDAPlace(0), ['X'], 'Out', check_prim=True, check_pir=True
         )
 
 
@@ -300,10 +300,10 @@ def setUp(self):
         self.outputs = {'Out': out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', numeric_grad_delta=0.5)
+        self.check_grad(['X'], 'Out', numeric_grad_delta=0.5, check_pir=True)
 
     def config(self):
         """
@@ -329,10 +329,10 @@ def setUp(self):
         self.outputs = {'Out': out}
 
     def test_check_output(self):
-        self.check_output()
+        self.check_output(check_pir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out')
+        self.check_grad(['X'], 'Out', check_pir=True)
 
     def config(self):
         """
@@ -603,11 +603,15 @@ def test_axis_maxsize():
 
 
 class TestCheckOutType(unittest.TestCase):
+    @test_with_pir_api
     def test_out_type(self):
         data = paddle.static.data(shape=[16, 10], dtype='int64', name='x')
         index = paddle.static.data(shape=[4], dtype='int64', name='index')
         out = paddle.gather(data, index)
-        self.assertTrue(out.dtype == core.VarDesc.VarType.INT64)
+        self.assertTrue(
+            out.dtype == core.VarDesc.VarType.INT64
+            or out.dtype == core.DataType.INT64
+        )
 
     def test_pir_out_type(self):
         with paddle.pir_utils.IrGuard():
