modify unsqueeze (#57681)

PaddlePaddle · Sep 25, 2023 · f53a032 · f53a032
1 parent 1c5abfa
commit f53a032
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 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
@@ -60,6 +60,8 @@
     'stack',
     'poisson',
     'gumbel_softmax',
+    'squeeze',
+    'unsqueeze',
     'tril',
     'triu',
 ]
@@ -100,6 +102,8 @@
     'stack',
     'poisson',
     'gumbel_softmax',
+    'squeeze',
+    'unsqueeze',
     'tril',
     'triu',
 ]
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
@@ -78,6 +78,8 @@
     'unsqueeze_grad',
     'poisson_grad',
     'gumbel_softmax_grad',
+    'squeeze_grad',
+    'unsqueeze_grad',
 ]
 
 

diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -2494,15 +2494,15 @@
   inplace : (out_grad -> x_grad)
 
 - backward_op : unsqueeze_double_grad
-  forward : unsqueeze_grad(Tensor xshape, Tensor grad_out, IntArray axes) -> Tensor(grad_x)
-  args : (Tensor grad_x_grad, IntArray axes)
+  forward : unsqueeze_grad(Tensor xshape, Tensor grad_out, IntArray axis) -> Tensor(grad_x)
+  args : (Tensor grad_x_grad, IntArray axis)
   output : Tensor(grad_out_grad), Tensor(xshape)
-  invoke : unsqueeze(grad_x_grad, axes)
+  invoke : unsqueeze(grad_x_grad, axis)
   intermediate : xshape
 
 - backward_op : unsqueeze_grad
-  forward : unsqueeze(Tensor x, IntArray axes) -> Tensor(out), Tensor(xshape)
-  args : (Tensor xshape, Tensor out_grad, IntArray axes)
+  forward : unsqueeze(Tensor x, IntArray axis) -> Tensor(out), Tensor(xshape)
+  args : (Tensor xshape, Tensor out_grad, IntArray axis)
   output : Tensor(x_grad)
   infer_meta :
     func : KernelWithXShapeInferMeta