add gather_nd (PaddlePaddle#57640)

* add gather_nd * Update paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py * Update test/legacy_test/test_gather_nd_op.py
Frida-a · Oct 14, 2023 · ba15a12 · ba15a12
1 parent 75678ea
commit ba15a12
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 16 deletions.
diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_gen_op_list.py
@@ -55,7 +55,9 @@
     'fused_softmax_mask_upper_triangle',
     'slice',
     'transpose',
-    'slice_double',
+    'slice_grad',
+    'gather_nd',
+    'stack',
     'poisson',
     'gumbel_softmax',
     'tril',
@@ -93,7 +95,9 @@
     'fused_softmax_mask_upper_triangle',
     'slice',
     'transpose',
-    'slice_double',
+    'slice_grad',
+    'gather_nd',
+    'stack',
     'poisson',
     'gumbel_softmax',
     'tril',

diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -72,6 +72,10 @@
     'layer_norm_grad',
     'embedding_grad',
     'scale_grad',
+    'gather_nd_grad',
+    'stack_grad',
+    'squeeze_grad',
+    'unsqueeze_grad',
     'poisson_grad',
     'gumbel_softmax_grad',
 ]
@@ -173,6 +177,10 @@
     'gumbel_softmax_grad',
     'split',
     'transpose',
+    'gather_nd_grad',
+    'stack_grad',
+    'squeeze_grad',
+    'unsqueeze_grad',
 ]
 
 

diff --git a/test/legacy_test/test_gather_nd_op.py b/test/legacy_test/test_gather_nd_op.py
@@ -56,7 +56,7 @@ def test_check_output(self):
         self.check_output(check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
 
 class TestGatherNdOpWithEmptyIndexFP16(TestGatherNdOpWithEmptyIndex):
@@ -80,7 +80,7 @@ def test_check_output(self):
     def test_check_grad(self):
         place = core.CUDAPlace(0)
         self.check_grad_with_place(
-            place, ['X'], 'Out', check_prim=True, check_new_ir=False
+            place, ['X'], 'Out', check_prim=True, check_new_ir=True
         )
 
 
@@ -117,7 +117,7 @@ def test_check_output(self):
         self.check_output(check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
 
 class TestGatherNdOpWithIndex1_ZeroDim(TestGatherNdOpWithIndex1):
@@ -168,7 +168,7 @@ def test_check_output(self):
     def test_check_grad(self):
         place = core.CUDAPlace(0)
         self.check_grad_with_place(
-            place, ['X'], 'Out', check_prim=True, check_new_ir=False
+            place, ['X'], 'Out', check_prim=True, check_new_ir=True
         )
 
 
@@ -205,7 +205,7 @@ def test_check_output(self):
         self.check_output(check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
 
 class TestGatherNdOpWithLowIndexFP16(TestGatherNdOpWithLowIndex):
@@ -233,7 +233,7 @@ def test_check_grad(self):
             ['X'],
             'Out',
             check_prim=True,
-            check_new_ir=False,
+            check_new_ir=True,
             numeric_grad_delta=0.5,
         )
 
@@ -280,7 +280,7 @@ def test_check_grad(self):
             ['X'],
             'Out',
             check_prim=True,
-            check_new_ir=False,
+            check_new_ir=True,
             numeric_grad_delta=0.05,
         )
 
@@ -310,7 +310,7 @@ def test_check_grad(self):
             ['X'],
             'Out',
             check_prim=True,
-            check_new_ir=False,
+            check_new_ir=True,
             numeric_grad_delta=0.5,
         )
 
@@ -345,7 +345,7 @@ def test_check_output(self):
         self.check_output(check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
 
 class TestGatherNdOpWithSameIndexAsXFP16(TestGatherNdOpWithSameIndexAsX):
@@ -373,7 +373,7 @@ def test_check_grad(self):
             ['X'],
             'Out',
             check_prim=True,
-            check_new_ir=False,
+            check_new_ir=True,
             numeric_grad_delta=0.5,
         )
 
@@ -410,7 +410,7 @@ def test_check_output(self):
         self.check_output(check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
 
 class TestGatherNdOpWithHighRankSameFP16(TestGatherNdOpWithHighRankSame):
@@ -434,7 +434,7 @@ def test_check_output(self):
     def test_check_grad(self):
         place = core.CUDAPlace(0)
         self.check_grad_with_place(
-            place, ['X'], 'Out', check_prim=True, check_new_ir=False
+            place, ['X'], 'Out', check_prim=True, check_new_ir=True
         )
 
 
@@ -471,7 +471,7 @@ def test_check_output(self):
         self.check_output(check_new_ir=True)
 
     def test_check_grad(self):
-        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=False)
+        self.check_grad(['X'], 'Out', check_prim=True, check_new_ir=True)
 
 
 class TestGatherNdOpWithHighRankDiffFP16(TestGatherNdOpWithHighRankDiff):
@@ -495,7 +495,7 @@ def test_check_output(self):
     def test_check_grad(self):
         place = core.CUDAPlace(0)
         self.check_grad_with_place(
-            place, ['X'], 'Out', check_prim=True, check_new_ir=False
+            place, ['X'], 'Out', check_prim=True, check_new_ir=True
         )