PaddlePaddle
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.cc
Copy file name to clipboard
-3Lines changed: 0 additions & 3 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/dynamic_reshape_pass.cc
Copy file name to clipboard
-3Lines changed: 0 additions & 3 deletions
diff --git a/‎paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
Copy file name to clipboard
+4-19Lines changed: 4 additions & 19 deletions b/‎paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
Copy file name to clipboard
+4-19Lines changed: 4 additions & 19 deletions
diff --git a/‎paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
Copy file name to clipboard
-2Lines changed: 0 additions & 2 deletions b/‎paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
Copy file name to clipboard
-2Lines changed: 0 additions & 2 deletions
diff --git a/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
Copy file name to clipboard
-2Lines changed: 0 additions & 2 deletions b/‎paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
Copy file name to clipboard
-2Lines changed: 0 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/transforms/gpu/fused_rotary_position_embedding.cc
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎paddle/fluid/pir/transforms/gpu/fused_rotary_position_embedding.cc
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
diff --git a/‎paddle/fluid/pir/transforms/onednn/operator_unsqueeze_onednn_fuse_pass.cc
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions b/‎paddle/fluid/pir/transforms/onednn/operator_unsqueeze_onednn_fuse_pass.cc
Copy file name to clipboard
+1-2Lines changed: 1 addition & 2 deletions
diff --git a/‎paddle/fluid/primitive/composite/composite.h
Copy file name to clipboard
+2-4Lines changed: 2 additions & 4 deletions b/‎paddle/fluid/primitive/composite/composite.h
Copy file name to clipboard
+2-4Lines changed: 2 additions & 4 deletions
diff --git a/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions b/‎paddle/fluid/primitive/rule/vjp/details.h
Copy file name to clipboard
+2-2Lines changed: 2 additions & 2 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/unsqueeze.cc
Copy file name to clipboard
+18-5Lines changed: 18 additions & 5 deletions b/‎paddle/phi/infermeta/spmd_rules/unsqueeze.cc
Copy file name to clipboard
+18-5Lines changed: 18 additions & 5 deletions
@@ -80,9 +80,6 @@ bool ReplaceOpWithReshapeOp(pir::Operation* op,
       op->operand_source(0), cinn_generate_shape.result(0));
 
   rewriter.ReplaceAllUsesWith(output, pd_reshape.result(0));
-  // if (with_xshape && op->num_results() == 2U) {
-  //   rewriter.ReplaceAllUsesWith(op->result(1), pd_reshape.result(1));
-  // }
   rewriter.EraseOp(op);
   return true;
 }
 
@@ -81,16 +81,13 @@ void ReplaceWithCinnReshapeOp(OpT op,
                               const std::vector<int> &out_shape) {
   PADDLE_ENFORCE_EQ(
       op->num_results(),
-      2U,
+      1U,
       ::common::errors::PreconditionNotMet(
-          "The size of source op outputs must be 2, but received %d.",
+          "The size of source op outputs must be 1, but received %d.",
           op->num_results()));
   auto cinn_reshape = rewriter.Build<cinn::dialect::ReshapeOp>(
       op->operand_source(0), out_shape);
-  auto generate_xshape =
-      rewriter.Build<cinn::dialect::GenerateXShapeOp>(op->operand_source(0));
   rewriter.ReplaceAllUsesWith(op.result(0), cinn_reshape.result(0));
-  rewriter.ReplaceAllUsesWith(op.result(1), generate_xshape.result(0));
 }
 
 }  // namespace
@@ -271,15 +268,7 @@ class ReshapeOpPattern
             out_shape_attr[i].dyn_cast<::pir::Int64Attribute>().data());
       }
     }
-    PADDLE_ENFORCE_EQ(
-        op->num_results(),
-        1U,
-        ::common::errors::PreconditionNotMet(
-            "The size of source op outputs must be 1, but received %d.",
-            op->num_results()));
-    auto cinn_reshape = rewriter.Build<cinn::dialect::ReshapeOp>(
-        op->operand_source(0), vec_out_shape);
-    rewriter.ReplaceAllUsesWith(op.result(0), cinn_reshape.result(0));
+    ReplaceWithCinnReshapeOp(op, rewriter, vec_out_shape);
     rewriter.EraseOp(op);
   }
 };
@@ -919,9 +908,7 @@ class SqueezeOpPattern
                   in_shape[i]));
         }
       }
-      auto cinn_reshape = rewriter.Build<cinn::dialect::ReshapeOp>(
-          op->operand_source(0), output_shape);
-      rewriter.ReplaceAllUsesWith(op.result(0), cinn_reshape.result(0));
+      ReplaceWithCinnReshapeOp(op, rewriter, output_shape);
       rewriter.EraseOp(op);
 
       return true;
@@ -1051,9 +1038,7 @@ class FlattenOpPattern
 
     auto reshape_op = rewriter.Build<paddle::dialect::ReshapeOp>(
         op->operand_source(0), new_shape);
-
     reshape_op.result(0).set_type(op.result(0).type());
-
     rewriter.ReplaceAllUsesWith(op.result(0), reshape_op.result(0));
 
     rewriter.EraseOp(op);
 
@@ -149,7 +149,6 @@
     "tensor_unfold",
     "transpose",
     "unbind",
-    "unsqueeze_infer",
     "unsqueeze",
     "view_shape",
     "view_dtype",
@@ -168,7 +167,6 @@
     "tensor_unfold_",
     "transpose_",
     "unbind_",
-    "unsqueeze_infer_",
     "view_shape_",
     "view_dtype_",
 }
 
@@ -3818,6 +3818,8 @@ OpTranslator::OpTranslator() {
       WithXShapeGradOpTranscriber<dialect::FlattenGradOp>();
   special_handlers["squeeze2_grad"] =
       WithXShapeAndAxisGradOpTranscriber<dialect::SqueezeGradOp>();
+  special_handlers["unsqueeze2_grad"] =
+      WithXShapeAndAxisGradOpTranscriber<dialect::UnsqueezeGradOp>();
 }
 }  // namespace translator
 }  // namespace paddle
@@ -2760,8 +2760,6 @@ bool UnsqueezeOpInferSymbolicShape(
 
   pir::Value res = op->result(0);
   infer_context->SetShapeOrDataForValue(res, shape_data);
-  infer_context->SetShapeOrDataForValue(
-      op->result(1), CreateShapeOrDataForXShape(x_shape_or_data));
 
   return true;
 }
 
@@ -110,7 +110,7 @@ class FusedRotaryPositionEmbeddingPattern : public paddle::drr::DrrPatternBase {
               {&pat.Tensor("squeeze_out_sin")});
 
     unsqueeze({&pat.Tensor("position_ids"), &full_11()},
-              {&pat.Tensor("unsqueeze_s_out_cos"), &pat.Tensor("xshape")});
+              {&pat.Tensor("unsqueeze_s_out_cos")});
 
     if (with_transpose_) {
       const auto &transpose_1 = pat.Op(paddle::dialect::TransposeOp::name(),
@@ -136,13 +136,13 @@ class FusedRotaryPositionEmbeddingPattern : public paddle::drr::DrrPatternBase {
     }
 
     unsqueeze_1({&pat.Tensor("gather_nd_out_cos"), &full_10()},
-                {&pat.Tensor("unsqueeze_out_cos"), &pat.Tensor("xshape")});
+                {&pat.Tensor("unsqueeze_out_cos")});
 
     unsqueeze_4({&pat.Tensor("position_ids"), &full_8()},
-                {&pat.Tensor("unsqueeze_s_out_sin"), &pat.Tensor("xshape")});
+                {&pat.Tensor("unsqueeze_s_out_sin")});
 
     unsqueeze_2({&pat.Tensor("gather_nd_out_sin"), &full_9()},
-                {&pat.Tensor("unsqueeze_out_sin"), &pat.Tensor("xshape")});
+                {&pat.Tensor("unsqueeze_out_sin")});
 
     pat.Tensor("tmp_25") =
         multiply1(pat.Tensor("q"), pat.Tensor("unsqueeze_out_cos"));
 
@@ -86,8 +86,7 @@ class OperatorUnsqueezeFusePattern : public paddle::drr::DrrPatternBase {
     const auto &full_1 = pat.Op(paddle::dialect::FullIntArrayOp::name(),
                                 {{"value", pat.Attr("full_1_value")}});
 
-    unsqueeze({&pat.Tensor("Out"), &full_1()},
-              {&pat.Tensor("Unsqueeze_out"), &pat.Tensor("Xshape")});
+    unsqueeze({&pat.Tensor("Out"), &full_1()}, {&pat.Tensor("Unsqueeze_out")});
 
     if (fusable_ops_ == paddle::onednn::dialect::FusedTransposeOp::name() ||
         fusable_ops_ ==
 
@@ -579,12 +579,10 @@ Tensor squeeze_decomp(const Tensor& x, const IntArray& axis) {
 }
 
 template <typename T>
-std::tuple<Tensor, Tensor> unsqueeze_decomp(const Tensor& x,
-                                            const IntArray& axis) {
+Tensor unsqueeze_decomp(const Tensor& x, const IntArray& axis) {
   auto out_shape = get_expand_dims(x, axis.GetData());
   Tensor out = reshape<T>(x, out_shape);
-  Tensor xshape;
-  return std::make_tuple(out, xshape);
+  return out;
 }
 
 template <typename T>
 
@@ -1130,7 +1130,7 @@ void squeeze_grad(const Tensor& x,
 }
 
 template <typename T>
-void unsqueeze_grad(const Tensor& xshape,
+void unsqueeze_grad(const Tensor& x,
                     const Tensor& out_grad,
                     const IntArray& axis,
                     Tensor* x_grad) {
@@ -1146,7 +1146,7 @@ void unsqueeze_grad(const Tensor& xshape,
     // for axis = [0, 3, 3], it outputs [0, 3, 3+1], because unsqueeze support
     // duplicated axis.
     std::vector<int64_t> output_axis;
-    const int64_t x_rank = xshape.dims().size() - 1;
+    const int64_t x_rank = x.dims().size();
     const std::vector<int64_t> axis_data = axis.GetData();
     for (size_t i = 0; i < axis_data.size(); ++i) {
       int64_t value = axis_data[i];
 
@@ -152,8 +152,7 @@ SpmdInfo UnsqueezeInferSpmd(const DistMetaTensor& x,
           << "]\n Out dims_mapping: [" << str_join(dims_mapping_vec[1])
           << "]\n\n";
 
-  return {{x_dist_attr_dst},
-          {out_dist_attr, CreateUnsqueezeXshape(x_dist_attr_dst)}};
+  return {{x_dist_attr_dst}, {out_dist_attr}};
 }
 
 SpmdInfo UnsqueezeInferSpmdReverse(const DistMetaTensor& x,
@@ -225,13 +224,27 @@ SpmdInfo UnsqueezeInferSpmdReverse(const DistMetaTensor& x,
           << "dims_mapping_dst: [" << str_join(dims_mapping_vec[0]) << "]";
   VLOG(4) << "X dims_mapping: [" << str_join(dims_mapping_vec[1]) << "]\n\n";
 
-  return {{x_dist_attr},
-          {out_dist_attr_dst, CreateUnsqueezeXshape(x_dist_attr)}};
+  return {{x_dist_attr}, {out_dist_attr_dst}};
 }
 
-SpmdInfo UnsqueezeGradInferSpmd(const DistMetaTensor& xshape,
+SpmdInfo UnsqueezeGradInferSpmd(const DistMetaTensor& x,
                                 const DistMetaTensor& out_grad,
                                 const IntArray& axis) {
+  auto shape = phi::vectorize(x.dims());
+  const auto& spmd = ReshapeInferSpmd(out_grad, shape);
+  return {{x.dist_attr(), spmd.first[0]}, {spmd.second[0]}};
+}
+
+SpmdInfo UnsqueezeWithXShapeInferSpmd(const DistMetaTensor& x,
+                                      const std::vector<int64_t>& axis) {
+  const auto& spmd = UnsqueezeInferSpmd(x, axis);
+  const auto& x_dist_attr = PADDLE_GET_CONST(TensorDistAttr, spmd.first[0]);
+  return {{x_dist_attr}, {spmd.second[0], CreateUnsqueezeXshape(x_dist_attr)}};
+}
+
+SpmdInfo UnsqueezeWithXShapeGradInferSpmd(const DistMetaTensor& xshape,
+                                          const DistMetaTensor& out_grad,
+                                          const IntArray& axis) {
   auto shape = phi::vectorize(xshape.dims());
   shape = std::vector<int64_t>(shape.begin() + 1, shape.end());
   const auto& spmd = ReshapeInferSpmd(out_grad, shape);
Original file line number	Diff line number	Diff line change
`@@ -80,9 +80,6 @@ bool ReplaceOpWithReshapeOp(pir::Operation* op,`
`80`	`80`	`op->operand_source(0), cinn_generate_shape.result(0));`
`81`	`81`
`82`	`82`	`rewriter.ReplaceAllUsesWith(output, pd_reshape.result(0));`
`83`		`- // if (with_xshape && op->num_results() == 2U) {`
`84`		`- // rewriter.ReplaceAllUsesWith(op->result(1), pd_reshape.result(1));`
`85`		`- // }`
`86`	`83`	`rewriter.EraseOp(op);`
`87`	`84`	`return true;`
`88`	`85`	`}`
Original file line number	Diff line number	Diff line change
`@@ -3818,6 +3818,8 @@ OpTranslator::OpTranslator() {`
`3818`	`3818`	`WithXShapeGradOpTranscriber<dialect::FlattenGradOp>();`
`3819`	`3819`	`special_handlers["squeeze2_grad"] =`
`3820`	`3820`	`WithXShapeAndAxisGradOpTranscriber<dialect::SqueezeGradOp>();`
	`3821`	`+ special_handlers["unsqueeze2_grad"] =`
	`3822`	`+ WithXShapeAndAxisGradOpTranscriber<dialect::UnsqueezeGradOp>();`
`3821`	`3823`	`}`
`3822`	`3824`	`} // namespace translator`
`3823`	`3825`	`} // namespace paddle`
Original file line number	Diff line number	Diff line change
`@@ -2760,8 +2760,6 @@ bool UnsqueezeOpInferSymbolicShape(`
`2760`	`2760`
`2761`	`2761`	`pir::Value res = op->result(0);`
`2762`	`2762`	`infer_context->SetShapeOrDataForValue(res, shape_data);`
`2763`		`- infer_context->SetShapeOrDataForValue(`
`2764`		`- op->result(1), CreateShapeOrDataForXShape(x_shape_or_data));`
`2765`	`2763`
`2766`	`2764`	`return true;`
`2767`	`2765`	`}`
Original file line number	Diff line number	Diff line change
`@@ -579,12 +579,10 @@ Tensor squeeze_decomp(const Tensor& x, const IntArray& axis) {`
`579`	`579`	`}`
`580`	`580`
`581`	`581`	`template <typename T>`
`582`		`-std::tuple<Tensor, Tensor> unsqueeze_decomp(const Tensor& x,`
`583`		`- const IntArray& axis) {`
	`582`	`+Tensor unsqueeze_decomp(const Tensor& x, const IntArray& axis) {`
`584`	`583`	`auto out_shape = get_expand_dims(x, axis.GetData());`
`585`	`584`	`Tensor out = reshape<T>(x, out_shape);`
`586`		`- Tensor xshape;`
`587`		`- return std::make_tuple(out, xshape);`
	`585`	`+ return out;`
`588`	`586`	`}`
`589`	`587`
`590`	`588`	`template <typename T>`