PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+156-3Lines changed: 156 additions & 3 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.cc
Copy file name to clipboard
+156-3Lines changed: 156 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+33-3Lines changed: 33 additions & 3 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_op.h
Copy file name to clipboard
+33-3Lines changed: 33 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
Copy file name to clipboard
+46-1Lines changed: 46 additions & 1 deletion b/‎paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
Copy file name to clipboard
+46-1Lines changed: 46 additions & 1 deletion
diff --git a/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions b/‎paddle/phi/infermeta/unary.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/phi/infermeta/unary.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
@@ -173,6 +173,17 @@ std::tuple<pir::OpResult, pir::OpResult> array_to_tensor(pir::Value x,
   return std::make_tuple(array_to_tensor.result(0), array_to_tensor.result(1));
 }
 
+pir::OpResult tensor_to_array(pir::Value x,
+                              pir::Value out_grad,
+                              int axis,
+                              bool use_stack) {
+  auto tensor_to_array = ApiBuilder::Instance()
+                             .GetBuilder()
+                             ->Build<paddle::dialect::TensorToArrayOp>(
+                                 x, out_grad, axis, use_stack);
+  return tensor_to_array.result(0);
+}
+
 pir::OpResult add_n_array(const std::vector<pir::Value>& inputs) {
   auto inputs_combine_op =
       ApiBuilder::Instance().GetBuilder()->Build<pir::CombineOp>(inputs);
 
@@ -74,6 +74,11 @@ std::tuple<pir::OpResult, pir::OpResult> array_to_tensor(pir::Value x,
                                                          int axis,
                                                          bool use_stack);
 
+pir::OpResult tensor_to_array(pir::Value x,
+                              pir::Value out_grad,
+                              int axis,
+                              bool use_stack);
+
 pir::OpResult add_n_array(const std::vector<pir::Value>& inputs);
 
 pir::OpResult slice_array_dense(pir::Value input, pir::Value starts);
 
@@ -22,9 +22,9 @@ paddle::dialect::AddNOp, paddle::dialect::AddN_Op,
     paddle::dialect::ArrayReadOp, paddle::dialect::ArrayWrite_Op,
     paddle::dialect::SliceArrayOp, paddle::dialect::SliceArrayDenseOp,
     paddle::dialect::AssignArray_Op, paddle::dialect::ArrayToTensorOp,
-    paddle::dialect::SelectInputOp, paddle::dialect::IncrementOp,
-    paddle::dialect::Increment_Op, paddle::dialect::ShapeBroadcastOp,
-    paddle::dialect::MemcpyD2hMultiIoOp
+    paddle::dialect::TensorToArrayOp, paddle::dialect::SelectInputOp,
+    paddle::dialect::IncrementOp, paddle::dialect::Increment_Op,
+    paddle::dialect::ShapeBroadcastOp, paddle::dialect::MemcpyD2hMultiIoOp
 #else
 
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
@@ -1085,6 +1085,7 @@ void SplitGradOp::Build(pir::Builder &builder,
       dense_x_grad.offset());
   argument_outputs.push_back(x_grad_dense_tensor_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void SplitGradOp::Build(pir::Builder &builder,
@@ -1142,6 +1143,7 @@ void SplitGradOp::Build(pir::Builder &builder,
       dense_x_grad.offset());
   argument_outputs.push_back(x_grad_dense_tensor_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void SplitGradOp::VerifySig() {
@@ -1250,6 +1252,7 @@ void CreateArrayOp::Build(pir::Builder &builder,
       dense_out.layout());
   argument_outputs.push_back(out_dense_tensor_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void CreateArrayOp::VerifySig() {
@@ -1885,6 +1888,7 @@ void ArrayToTensorOp::Build(pir::Builder &builder,             // NOLINT
       dense_out_index.offset());
   argument_outputs.push_back(out_index_dense_tensor_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void ArrayToTensorOp::VerifySig() {
@@ -1940,6 +1944,153 @@ void ArrayToTensorOp::InferMeta(phi::InferMetaContext *infer_meta) {
   fn(infer_meta);
 }
 
+const char *TensorToArrayOp::attributes_name[2] = {"axis", "use_stack"};
+
+OpInfoTuple TensorToArrayOp::GetOpInfo() {
+  std::vector<paddle::dialect::OpInputInfo> inputs = {
+      paddle::dialect::OpInputInfo("x",
+                                   "paddle::dialect::DenseTensorArrayType",
+                                   false,
+                                   false,
+                                   false,
+                                   true),
+      paddle::dialect::OpInputInfo("out_grad",
+                                   "paddle::dialect::DenseTensorType",
+                                   false,
+                                   false,
+                                   false,
+                                   true)};
+
+  std::vector<paddle::dialect::OpAttributeInfo> attributes = {
+      paddle::dialect::OpAttributeInfo("axis", "pir::Int32Attribute", ""),
+      paddle::dialect::OpAttributeInfo("use_stack", "pir::BoolAttribute", "")};
+
+  std::vector<paddle::dialect::OpOutputInfo> outputs = {
+      paddle::dialect::OpOutputInfo(
+          "out", "paddle::dialect::DenseTensorArrayType", false, false)};
+
+  paddle::dialect::OpRunTimeInfo run_time_info =
+      paddle::dialect::OpRunTimeInfo("TensorToArrayInferMeta",
+                                     {"x", "axis", "use_stack"},
+                                     "tensor_to_array",
+                                     {"x", "axis", "use_stack"},
+                                     {"x"},
+                                     {},
+                                     {},
+                                     {});
+  return std::make_tuple(
+      inputs, attributes, outputs, run_time_info, "tensor_to_array");
+}
+
+void TensorToArrayOp::Build(pir::Builder &builder,             // NOLINT
+                            pir::OperationArgument &argument,  // NOLINT
+                            pir::Value x_,
+                            pir::Value out_grad_,
+                            int axis,
+                            bool use_stack) {
+  VLOG(4) << "Start build TensorToArrayOp";
+  VLOG(4) << "Builder construction inputs";
+  argument.AddInputs({x_, out_grad_});
+
+  VLOG(4) << "Builder construction attributes";
+  pir::Attribute attr_axis =
+      pir::Int32Attribute::get(pir::IrContext::Instance(), axis);
+  argument.AddAttribute("axis", attr_axis);
+  pir::Attribute attr_use_stack =
+      pir::BoolAttribute::get(pir::IrContext::Instance(), use_stack);
+  argument.AddAttribute("use_stack", attr_use_stack);
+
+  VLOG(4) << "Builder construction outputs";
+  paddle::dialect::DenseTensorArrayType x =
+      x_.type().dyn_cast<paddle::dialect::DenseTensorArrayType>();
+  paddle::dialect::IrTensor dense_x(
+      paddle::dialect::TransToPhiDataType(x.dtype()), {}, x.data_layout(), {});
+
+  paddle::dialect::DenseTensorType out_grad =
+      out_grad_.type().dyn_cast<paddle::dialect::DenseTensorType>();
+  paddle::dialect::IrTensor dense_out_grad(
+      paddle::dialect::TransToPhiDataType(out_grad.dtype()),
+      out_grad.dims(),
+      out_grad.data_layout(),
+      out_grad.lod(),
+      out_grad.offset());
+
+  VLOG(4) << "Builder construction  meta_x, meta_out_grad";
+  paddle::dialect::IrMetaTensor meta_out_grad(&dense_out_grad);
+  paddle::dialect::IrMetaTensor meta_x(&dense_x);
+
+  paddle::dialect::IrTensor dense_x_grad;
+  paddle::dialect::IrMetaTensor meta_x_grad(&dense_x_grad);
+
+  phi::TensorToArrayInferMeta(
+      meta_x, meta_out_grad, axis, use_stack, &meta_x_grad);
+
+  std::vector<pir::Type> argument_outputs;
+  pir::Type out_dense_tensor_array_type =
+      paddle::dialect::DenseTensorArrayType::get(
+          pir::IrContext::Instance(),
+          paddle::dialect::TransToIrDataType(dense_x_grad.dtype()),
+          dense_x_grad.layout());
+  argument_outputs.push_back(out_dense_tensor_array_type);
+  argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
+}
+
+void TensorToArrayOp::VerifySig() {
+  VLOG(4) << "Start Verifying inputs, outputs and attributes for: "
+             "TensorToArrayOp.";
+  VLOG(4) << "Verifying inputs:";
+  {
+    auto input_size = num_operands();
+    PADDLE_ENFORCE_EQ(
+        input_size,
+        2u,
+        phi::errors::PreconditionNotMet(
+            "The size %d of inputs must be equal to 2.", input_size));
+
+    PADDLE_ENFORCE((*this)
+                       ->operand_source(0)
+                       .type()
+                       .isa<paddle::dialect::DenseTensorArrayType>(),
+                   phi::errors::PreconditionNotMet(
+                       "Type validation failed for the 0th input."));
+    PADDLE_ENFORCE((*this)
+                       ->operand_source(1)
+                       .type()
+                       .isa<paddle::dialect::DenseTensorType>(),
+                   phi::errors::PreconditionNotMet(
+                       "Type validation failed for the 1th input."));
+  }
+
+  VLOG(4) << "Verifying attributes:";
+  {
+    auto &attributes = this->attributes();
+    PADDLE_ENFORCE(attributes.count("axis") > 0, "axis does not exist.");
+    PADDLE_ENFORCE(attributes.count("use_stack") > 0,
+                   "use_stack does not exist.");
+  }
+
+  VLOG(4) << "Verifying outputs:";
+  {
+    auto output_size = num_results();
+    PADDLE_ENFORCE_EQ(
+        output_size,
+        1u,
+        phi::errors::PreconditionNotMet(
+            "The size %d of outputs must be equal to 1.", output_size));
+    PADDLE_ENFORCE(
+        (*this)->result(0).type().isa<paddle::dialect::DenseTensorArrayType>(),
+        phi::errors::PreconditionNotMet(
+            "Type validation failed for the 0th output."));
+  }
+  VLOG(4) << "End Verifying for: TensorToArrayOp.";
+}
+
+void TensorToArrayOp::InferMeta(phi::InferMetaContext *infer_meta) {
+  auto fn = PD_INFER_META(phi::TensorToArrayInferMeta);
+  fn(infer_meta);
+}
+
 const char *SliceArrayOp::attributes_name[2] = {"starts", "ends"};
 
 OpInfoTuple SliceArrayOp::GetOpInfo() {
@@ -2161,6 +2312,7 @@ void SliceArrayDenseOp::Build(pir::Builder &builder,             // NOLINT
       dense_out.offset());
   argument_outputs.push_back(out_dense_tensor_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void SliceArrayDenseOp::InferMeta(phi::InferMetaContext *infer_meta) {
@@ -3136,6 +3288,7 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SliceArrayOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SliceArrayDenseOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AssignArray_Op)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayToTensorOp)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::TensorToArrayOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 
@@ -316,8 +316,10 @@ class ArrayWrite_Op : public pir::Op<ArrayWrite_Op,
       const std::vector<std::vector<bool>> &stop_gradients);
 };
 
-class ArrayToTensorOp
-    : public pir::Op<ArrayToTensorOp, OpYamlInfoInterface, InferMetaInterface> {
+class ArrayToTensorOp : public pir::Op<ArrayToTensorOp,
+                                       OpYamlInfoInterface,
+                                       paddle::dialect::VjpInterface,
+                                       InferMetaInterface> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_to_tensor"; }
@@ -332,7 +334,34 @@ class ArrayToTensorOp
   void VerifySig();
   pir::Value x() { return operand_source(0); }
   pir::OpResult out() { return result(0); }
-  pir::OpResult out_index() { return result(2); }
+  pir::OpResult out_index() { return result(1); }
+  static void InferMeta(phi::InferMetaContext *infer_meta);
+  static std::vector<std::vector<pir::OpResult>> Vjp(
+      pir::Operation *op,
+      const std::vector<std::vector<pir::Value>> &inputs_,
+      const std::vector<std::vector<pir::Value>> &outputs,
+      const std::vector<std::vector<pir::Value>> &out_grads,
+      const std::vector<std::vector<bool>> &stop_gradients);
+};
+
+class TensorToArrayOp
+    : public pir::Op<TensorToArrayOp, OpYamlInfoInterface, InferMetaInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "pd_op.tensor_to_array"; }
+  static constexpr uint32_t attributes_num = 2;
+  static const char *attributes_name[attributes_num];
+  static OpInfoTuple GetOpInfo();
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    pir::Value x,
+                    pir::Value out_grad,
+                    int axis,
+                    bool use_stack);
+  void VerifySig();
+  pir::Value x() { return operand_source(0); }
+  pir::Value out_grad() { return operand_source(1); }
+  pir::OpResult x_grad() { return result(0); }
   static void InferMeta(phi::InferMetaContext *infer_meta);
 };
 
@@ -630,6 +659,7 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SliceArrayOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SliceArrayDenseOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AssignArray_Op)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ArrayToTensorOp)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::TensorToArrayOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::ExpandOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IncrementOp)
 
@@ -256,7 +256,7 @@ std::vector<std::vector<pir::OpResult>> ArrayReadOp::Vjp(
       2,
       platform::errors::InvalidArgument(
           "Array_read op's outputs size should be 1, but now is %d.",
-          outputs.size()));
+          out_grads.size()));
 
   VLOG(6) << "Vjp prepare call  Array_read's vjp inteface";
 
@@ -269,5 +269,50 @@ std::vector<std::vector<pir::OpResult>> ArrayReadOp::Vjp(
   std::vector<std::vector<pir::OpResult>> res;
   return res;
 }
+
+std::vector<std::vector<pir::OpResult>> ArrayToTensorOp::Vjp(
+    pir::Operation* op,
+    const std::vector<std::vector<pir::Value>>& inputs_,
+    const std::vector<std::vector<pir::Value>>& outputs,
+    const std::vector<std::vector<pir::Value>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  PADDLE_ENFORCE_EQ(
+      inputs_.size(),
+      1,
+      platform::errors::InvalidArgument(
+          "Array_read op's inputs size should be 1, but now is %d.",
+          inputs_.size()));
+  PADDLE_ENFORCE_EQ(
+      outputs.size(),
+      2,
+      platform::errors::InvalidArgument(
+          "Array_read op's outputs size should be 2, but now is %d.",
+          outputs.size()));
+
+  PADDLE_ENFORCE_EQ(
+      out_grads.size(),
+      2,
+      platform::errors::InvalidArgument(
+          "Array_read op's outputs size should be 2, but now is %d.",
+          out_grads.size()));
+
+  VLOG(6) << "Vjp prepare Prepare attributes of array_to_tensor_grad";
+  int axis = op->attribute("axis").dyn_cast<pir::Int32Attribute>().data();
+  bool use_stack =
+      op->attribute("use_stack").dyn_cast<pir::BoolAttribute>().data();
+
+  VLOG(6) << "Vjp prepare call ArrayToTensor's vjp inteface";
+
+  pir::OpResult tensor_res = paddle::dialect::tensor_to_array(
+      inputs_[0][0], out_grads[0][0], axis, use_stack);
+
+  std::vector<std::vector<pir::OpResult>> res(1);
+  res[0].resize(1);
+  if (!stop_gradients[0][0]) {
+    res[0][0] = tensor_res;
+  }
+  return res;
+}
+
 }  // namespace dialect
 }  // namespace paddle
@@ -175,6 +175,15 @@ void ArrayToTensorInferMeta(const MetaTensor& x,
   out_index->set_dims(common::make_ddim({-1}));
 }
 
+void TensorToArrayInferMeta(const MetaTensor& x,
+                            const MetaTensor& out_grad,
+                            int axis,
+                            bool use_stack,
+                            MetaTensor* x_grad) {
+  x_grad->set_dtype(x.dtype());
+  x_grad->set_layout(x.layout());
+}
+
 void ArgMinMaxInferMeta(const MetaTensor& x,
                         const Scalar& axis,
                         bool keepdims,
 
@@ -68,6 +68,12 @@ void ArrayToTensorInferMeta(const MetaTensor& x,
                             MetaTensor* out_index,
                             MetaConfig config = MetaConfig());
 
+void TensorToArrayInferMeta(const MetaTensor& x,
+                            const MetaTensor& out_grad,
+                            int axis,
+                            bool use_stack,
+                            MetaTensor* x_grad);
+
 void AsRealInferMeta(const MetaTensor& input, MetaTensor* output);
 
 void AsComplexInferMeta(const MetaTensor& input, MetaTensor* output);