refine

oneflow/ir/include/OneFlow/OneFlowBase.td

@@ -183,7 +183,6 @@ class OneFlow_TFPoolGradBaseOp<string mnemonic, list<Trait> traits = []> :
   let summary = "OneFlow pooling grad operation, align with TensorFlow";
   let input = (ins
     AnyType:$x,
-    AnyType:$y,
     AnyType:$dy
   );
   let output = (outs AnyType:$dx);
@@ -257,7 +256,6 @@ class OneFlow_MaxPoolGradBaseOp<string mnemonic, list<Trait> traits = []> :
   let summary = "OneFlow Max Pooling Grad operation";
   let input = (ins
     AnyType:$x,
-    AnyType:$y,
     AnyType:$indice,
     AnyType:$dy
   );
@@ -284,7 +282,6 @@ class OneFlow_AvgPoolGradBaseOp<string mnemonic, list<Trait> traits = []> :
   let summary = "OneFlow Average Pooling Grad operation";
   let input = (ins
     AnyType:$x,
-    AnyType:$y,
     AnyType:$dy
   );
   let output = (outs

oneflow/user/ops/pool_op.cpp

@@ -93,6 +93,7 @@ GenBackwardOpConfFn MakeGenBackwardOpConfFn(const std::string& mode, const int32
       user_op::UserOpConfWrapperBuilder builder(op.op_name() + "_grad");
       user_op::UserOpConfWrapper grad_op =
           builder.Op(mode + "_pool_" + std::to_string(dim) + "d_grad")
+              .Input("x", op.input("x", 0))
               .Input("dy", op.GetGradTensorWithOpOutput("y", 0))
               .Output("dx")
               .Attr("data_format", op.attr<std::string>("data_format"))

oneflow/user/ops/pooling_op.cpp

@@ -143,6 +143,7 @@ GenBackwardOpConfFn MaxPoolMakeBackwardOpConfFn(const std::string& mode, const i
       user_op::UserOpConfWrapperBuilder builder(op.op_name() + "_grad");
       user_op::UserOpConfWrapper grad_op =
           builder.Op(mode + "pool_" + std::to_string(dim) + "d_grad")
+              .Input("x", op.input("x", 0))
               .Input("indice", op.output("indice", 0))
               .Input("dy", op.GetGradTensorWithOpOutput("y", 0))
               .Output("dx")
@@ -167,6 +168,7 @@ GenBackwardOpConfFn AvgPoolMakeBackwardOpConfFn(const int32_t dim) {
       user_op::UserOpConfWrapperBuilder builder(op.op_name() + "_grad");
       user_op::UserOpConfWrapper grad_op =
           builder.Op("avgpool_" + std::to_string(dim) + "d_grad")
+              .Input("x", op.input("x", 0))
               .Input("dy", op.GetGradTensorWithOpOutput("y", 0))
               .Output("dx")
               .Attr("data_format", op.attr<std::string>("data_format"))