Add yaml for flatten_contiguous_range OP

paddle/phi/kernels/flatten_grad_kernel.cc

@@ -21,8 +21,8 @@ namespace phi {
 
 template <typename T, typename Context>
 void FlattenGradKernel(const Context& dev_ctx,
-                       const DenseTensor& out_grad,
                        const DenseTensor& xshape,
+                       const DenseTensor& out_grad,
                        DenseTensor* x_grad) {
   auto xshape_dims = xshape.dims();
   dev_ctx.Alloc(x_grad, out_grad.dtype());

paddle/phi/kernels/flatten_grad_kernel.h

@@ -20,8 +20,8 @@ namespace phi {
 
 template <typename T, typename Context>
 void FlattenGradKernel(const Context& dev_ctx,
-                       const DenseTensor& out_grad,
                        const DenseTensor& xshape,
+                       const DenseTensor& out_grad,
                        DenseTensor* x_grad);
 
 }  // namespace phi

paddle/phi/ops/compat/flatten_sig.cc

@@ -31,7 +31,7 @@ KernelSignature FlattenOpArgumentMapping(const ArgumentMappingContext& ctx) {
 KernelSignature FlattenGradOpArgumentMapping(
     const ArgumentMappingContext& ctx) {
   return KernelSignature(
-      "flatten_grad", {GradVarName("Out"), "XShape"}, {}, {GradVarName("X")});
+      "flatten_grad", {"XShape", GradVarName("Out")}, {}, {GradVarName("X")});
 }
 
 }  // namespace phi

paddle/phi/tests/api/CMakeLists.txt

@@ -12,7 +12,6 @@ cc_test(test_dot_api SRCS test_dot_api.cc DEPS ${COMMON_API_TEST_DEPS})
 cc_test(test_matmul_api SRCS test_matmul_api.cc DEPS ${COMMON_API_TEST_DEPS})
 cc_test(test_empty_api SRCS test_empty_api.cc DEPS ${COMMON_API_TEST_DEPS})
 cc_test(test_fill_api SRCS test_fill_api.cc DEPS ${COMMON_API_TEST_DEPS})
-cc_test(test_flatten_api SRCS test_flatten_api.cc DEPS ${COMMON_API_TEST_DEPS})
 cc_test(test_elementwise_api SRCS test_elementwise_api.cc DEPS ${COMMON_API_TEST_DEPS})
 cc_test(test_cast_api SRCS test_cast_api.cc DEPS ${COMMON_API_TEST_DEPS})
 cc_test(test_reshape_api SRCS test_reshape_api.cc DEPS ${COMMON_API_TEST_DEPS})

python/paddle/fluid/tests/unittests/test_flatten_contiguous_range_op.py

@@ -23,6 +23,8 @@
 
 class TestFlattenOp(OpTest):
     def setUp(self):
+        self.python_api = paddle.flatten
+        self.python_out_sig = ["Out"]
         self.op_type = "flatten_contiguous_range"
         self.start_axis = 0
         self.stop_axis = -1
@@ -35,10 +37,10 @@ def setUp(self):
         }
 
     def test_check_output(self):
-        self.check_output(no_check_set=["XShape"])
+        self.check_output(no_check_set=["XShape"], check_eager=True)
 
     def test_check_grad(self):
-        self.check_grad(["X"], "Out")
+        self.check_grad(["X"], "Out", check_eager=True)
 
     def init_test_case(self):
         self.in_shape = (3, 2, 5, 4)

python/paddle/tensor/manipulation.py

@@ -697,7 +697,11 @@ def flatten(x, start_axis=0, stop_axis=-1, name=None):
     if start_axis > stop_axis:
         raise ValueError("The stop_axis should be larger than stat_axis")
 
-    if paddle.in_dynamic_mode():
+    if in_dygraph_mode():
+        dy_out, _ = _C_ops.final_state_flatten(x, start_axis, stop_axis)
+        return dy_out
+
+    if _in_legacy_dygraph():
         dy_out, _ = _C_ops.flatten_contiguous_range(x, 'start_axis', start_axis,
                                                     'stop_axis', stop_axis)
         return dy_out

python/paddle/utils/code_gen/api.yaml

@@ -547,11 +547,15 @@
 
 - api : flatten
   args : (Tensor x, int start_axis, int stop_axis)
-  output : Tensor
+  output : Tensor(out), Tensor(xshape)
   infer_meta :
-    func : FlattenInferMeta
+    func : FlattenWithXShapeInferMeta
   kernel :
-    func : flatten
+    func : flatten_with_xshape
+    backend : x
+  inplace : (x -> out)
+  view : (x -> out)
+  backward : flatten_grad
 
 # flip
 - api : flip

python/paddle/utils/code_gen/backward.yaml

@@ -349,6 +349,19 @@
   kernel :
     func : expm1_grad
 
+- backward_api : flatten_grad
+  forward : flatten(Tensor x, int start_axis, int stop_axis) -> Tensor(out), Tensor(xshape)
+  args : (Tensor xshape, Tensor out_grad)
+  output : Tensor(x_grad)
+  infer_meta :
+    func :  KernelWithXShapeInferMeta
+    param : [xshape]
+  kernel :
+    func : flatten_grad
+    data_type: out_grad
+    backend: out_grad
+    layout: out_grad
+
 - backward_api : floor_grad
   forward : floor(Tensor x) -> Tensor(out)
   args : (Tensor out_grad)

tools/infrt/skipped_phi_api.json

@@ -1,4 +1,4 @@
 {
-"phi_apis":["conj", "nll_loss"],
+"phi_apis":["conj", "nll_loss", "flatten"],
 "phi_kernels":["equal_all"]
 }