PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
Copy file name to clipboard
+32-27Lines changed: 32 additions & 27 deletions b/‎paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
Copy file name to clipboard
+32-27Lines changed: 32 additions & 27 deletions
diff --git a/‎python/paddle/autograd/backward_utils.py
Copy file name to clipboard
+18-2Lines changed: 18 additions & 2 deletions b/‎python/paddle/autograd/backward_utils.py
Copy file name to clipboard
+18-2Lines changed: 18 additions & 2 deletions
@@ -51,6 +51,7 @@ void IfOp::Build(pir::Builder &builder,             // NOLINT
   if (true_block && !true_block->empty() &&
       true_block->back().isa<pir::YieldOp>()) {
     auto &op = true_block->back();
+
     for (size_t i = 0; i < op.num_operands(); ++i) {
       argument.AddOutput(op.operand(i).type());
     }
@@ -192,13 +193,14 @@ std::vector<std::vector<pir::OpResult>> IfOp::Vjp(
     const std::vector<std::vector<pir::OpResult>> &outputs,
     const std::vector<std::vector<pir::Value>> &out_grads,
     const std::vector<std::vector<bool>> &stop_gradients) {
-  PADDLE_ENFORCE_EQ(inputs_.size() == 1u && inputs_[0].size() >= 1u,
-                    true,
-                    phi::errors::InvalidArgument(
-                        "if op's inputs' size should be 1, and the inputs[0] "
-                        "should be non-empty. "
-                        "Now the inputs's size is %d or inputs[0] is empty.",
-                        inputs_.size()));
+  PADDLE_ENFORCE_EQ(
+      inputs_.size() >= 1u,
+      true,
+      phi::errors::InvalidArgument("if op's inputs' size should greater_equal "
+                                   "to 1, and all the inputs[i] "
+                                   "should be 1 size. "
+                                   "Now the inputs's size is %d .",
+                                   inputs_.size()));
 
   VLOG(6) << "Prepare inputs for if_grad";
   auto cond_val = inputs_[0][0];
@@ -207,21 +209,20 @@ std::vector<std::vector<pir::OpResult>> IfOp::Vjp(
   VLOG(6) << "Prepare outputs for if_grad";
 
   std::vector<pir::Type> output_types;
-  for (size_t i = 0; i < inputs_[0].size(); ++i) {
-    if (!stop_gradients[0][i]) {
-      output_types.push_back(inputs_[0][i].type());
+  for (size_t i = 0; i < inputs_.size(); ++i) {
+    if (!stop_gradients[i][0]) {
+      output_types.push_back(inputs_[i][0].type());
     }
   }
 
   auto if_grad = ApiBuilder::Instance().GetBuilder()->Build<IfOp>(
       cond_val, std::move(output_types));
 
-  std::vector<std::vector<pir::OpResult>> res{
-      std::vector<pir::OpResult>(inputs_[0].size())};
-
-  for (size_t i = 0, j = 0; i < inputs_[0].size(); ++i) {
-    if (!stop_gradients[0][i]) {
-      res[0][i] = if_grad->result(j++);
+  std::vector<std::vector<pir::OpResult>> res{inputs_.size()};
+  for (size_t i = 0, j = 0; i < inputs_.size(); ++i) {
+    res[i].resize(1);
+    if (!stop_gradients[i][0]) {
+      res[i][0] = if_grad->result(j++);
     }
   }
   return res;
@@ -278,19 +279,23 @@ std::vector<std::vector<pir::OpResult>> TuplePushOpVjpInterfaceModel::Vjp(
     const std::vector<std::vector<pir::OpResult>> &outputs,
     const std::vector<std::vector<pir::Value>> &out_grads,
     const std::vector<std::vector<bool>> &stop_gradients) {
-  PADDLE_ENFORCE_EQ(inputs.size() == 1u && inputs[0].size() >= 1u,
-                    true,
-                    phi::errors::InvalidArgument(
-                        "tupe_push op's inputs' size should be 1, and the "
-                        "inputs[0] should be non-empty. "
-                        "Now the inputs's size is %d or inputs[0] is empty.",
-                        inputs.size()));
+  PADDLE_ENFORCE_EQ(
+      inputs.size() >= 1u,
+      true,
+      phi::errors::InvalidArgument(
+          "tupe_push op's inputs' size should be greater_equal than 1, and the "
+          "inputs[i] should be non-empty. "
+          "Now the inputs's size is %d.",
+          inputs.size()));
   auto pop_op = ApiBuilder::Instance().GetBuilder()->Build<TuplePopOp>(
       TuplePushOp::dyn_cast(op).outlet());
-  std::vector<std::vector<pir::OpResult>> res{
-      std::vector<pir::OpResult>{nullptr}};
-  for (size_t i = 0u; i < pop_op.num_results(); ++i) {
-    res[0].push_back(pop_op.result(i));
+  std::vector<std::vector<pir::OpResult>> res{inputs.size()};
+  res[0].resize(1);
+  for (size_t i = 1u; i < inputs.size(); ++i) {
+    res[i].resize(1);
+    if (!stop_gradients[i][0]) {
+      res[i][0] = pop_op.result(i - 1);
+    }
   }
   return res;
 }
 
@@ -22,8 +22,8 @@ class State:
 
     """
 
-    def __init__(self, program):
-        self.program = program
+    def __init__(self, block):
+        self.block = block
         # opresult -> list(list(opresult))
         self.value_to_valuegrad = collections.defaultdict(list)
         self.value_to_sumvaluegrad = collections.defaultdict(list)
@@ -52,3 +52,19 @@ def turn_map(self) -> None:
         for k, v in self.op_to_opgrad.items():
             if v != []:
                 self.opgrad_to_op[v[0]] = [k]
+
+    def copy(self, new_block):
+        state = State(new_block)
+        state.value_to_valuegrad = self.value_to_valuegrad.copy()
+        state.value_to_sumvaluegrad = self.value_to_sumvaluegrad.copy()
+
+        # operation -> list(operation)
+        state.op_to_opgrad = self.op_to_opgrad.copy()
+
+        # opresult -> list(opresult)
+        state.valuegrad_to_value = self.valuegrad_to_value.copy()
+        state.sumvaluegrad_to_value = self.sumvaluegrad_to_value.copy()
+        # operation -> list(operation)
+        state.opgrad_to_op = self.opgrad_to_op.copy()
+
+        return state