static graph autogen code for check_finite_and_unscale_ op (#54145)

* static graph autogen code for check_finite_and_unscale_ op

* bug fixed

paddle/fluid/operators/amp/check_finite_and_unscale_op.cc

-/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License. */
-
-#include "paddle/fluid/framework/infershape_utils.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/phi/core/infermeta_utils.h"
-#include "paddle/phi/infermeta/multiary.h"
-
-namespace paddle {
-namespace operators {
-
-class CheckFiniteAndUnscaleOp : public framework::OperatorWithKernel {
- public:
-  CheckFiniteAndUnscaleOp(const std::string& type,
-                          const framework::VariableNameMap& inputs,
-                          const framework::VariableNameMap& outputs,
-                          const framework::AttributeMap& attrs)
-      : OperatorWithKernel(type, inputs, outputs, attrs) {}
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    auto dtype = framework::proto::VarType::FP32;
-    if (ctx.MultiInputVar("X").size() >= 1) {
-      dtype = OperatorWithKernel::IndicateVarDataType(ctx, "X");
-    }
-    return phi::KernelKey(dtype, ctx.GetPlace());
-  }
-};
-
-class CheckFiniteAndUnscaleOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput(
-        "X",
-        "(Tensors) The input tensors of check_finite_and_unscale operator.")
-        .AsDuplicable();
-    AddInput("Scale",
-             "(Tensor) 1-dim tensor, the scale of check_finite_and_unscale "
-             "operator.");
-    AddOutput("Out",
-              "(Tensors) The scaled output tensor of "
-              "check_finite_and_unscale operator.")
-        .AsDuplicable();
-    AddOutput("FoundInfinite",
-              "(Tensor) 1-dim tensor, contains a bool scalar, which indicates "
-              "if there there is infinite or nan item in input X.");
-    AddComment(R"DOC(
-check_finite_and_unscale operator.
-Check if input X contains all finite data, if yes, scale it by input Scale.
-
-$$Out = X / scale$$
-
-If any tensor in X contains Inf or Nan, the Out will generate a indicator.
-FoundInfinite will be 1 (True), and Out will not be scaled. In this case, the data of
-Out should not be used, and its data may not be deterministic.
-Otherwise, FoundInfinite will be 0 (False).
-
-)DOC");
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-
-DECLARE_INFER_SHAPE_FUNCTOR(check_finite_and_unscale,
-                            CheckFiniteAndUnscaleInferShapeFunctor,
-                            PD_INFER_META(phi::CheckFiniteAndUnscaleInferMeta));
-REGISTER_OPERATOR(
-    check_finite_and_unscale,
-    ops::CheckFiniteAndUnscaleOp,
-    ops::CheckFiniteAndUnscaleOpMaker,
-    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
-    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
-    CheckFiniteAndUnscaleInferShapeFunctor);

paddle/fluid/operators/generator/get_expected_kernel_func.cc

@@ -61,6 +61,16 @@ static bool ReduceOpHasOptimizedOneDNNKernel(
   return true;
 }
 
+phi::KernelKey GetCheckFiniteAndUnscaleExpectedKernelType(
+    const framework::ExecutionContext& ctx,
+    const framework::OperatorWithKernel* op_ptr) {
+  auto dtype = framework::proto::VarType::FP32;
+  if (ctx.MultiInputVar("X").size() >= 1) {
+    dtype = op_ptr->IndicateVarDataType(ctx, "X");
+  }
+  return phi::KernelKey(dtype, ctx.GetPlace());
+}
+
 phi::KernelKey GetReduceExpectedKernelType(
     const framework::ExecutionContext& ctx,
     const framework::OperatorWithKernel* op_ptr) {

paddle/fluid/operators/generator/get_expected_kernel_func.h

@@ -24,6 +24,10 @@ phi::KernelKey GetReduceExpectedKernelType(
     const framework::ExecutionContext& ctx,
     const framework::OperatorWithKernel* op_ptr);
 
+phi::KernelKey GetCheckFiniteAndUnscaleExpectedKernelType(
+    const framework::ExecutionContext& ctx,
+    const framework::OperatorWithKernel* op_ptr);
+
 phi::KernelKey GetReduceGradExpectedKernelType(
     const framework::ExecutionContext& ctx,
     const framework::OperatorWithKernel* op_ptr);

paddle/phi/api/yaml/legacy_ops.yaml

@@ -150,18 +150,6 @@
     func : channel_shuffle
   backward : channel_shuffle_grad
 
-- op : check_finite_and_unscale_
-  args : (Tensor[] x, Tensor scale, Tensor input_found_infinite)
-  output : Tensor[](out){x.size()}, Tensor(output_found_infinite)
-  infer_meta :
-    func : CheckFiniteAndUnscaleInferMeta
-    param : [x, scale]
-  kernel :
-    func : check_finite_and_unscale
-    param : [x, scale]
-    data_type : x
-  inplace : (x -> out), (input_found_infinite -> output_found_infinite)
-
 - op : concat
   args : (Tensor[] x, Scalar(int64_t) axis)
   output : Tensor

paddle/phi/api/yaml/op_compat.yaml

@@ -415,6 +415,14 @@
   outputs :
     out : Out
 
+- op : check_finite_and_unscale_
+  inputs :
+    {x : X, scale: Scale}
+  outputs :
+    {out : Out, found_infinite: FoundInfinite}
+  get_expected_kernel_type :
+    check_finite_and_unscale_ : GetCheckFiniteAndUnscaleExpectedKernelType
+
 - op : cholesky
   inputs :
     x : X

paddle/phi/api/yaml/ops.yaml

@@ -378,6 +378,18 @@
     func : celu
   backward : celu_grad
 
+- op : check_finite_and_unscale_
+  args : (Tensor[] x, Tensor scale)
+  output : Tensor[](out){x.size()}, Tensor(found_infinite)
+  infer_meta :
+    func : CheckFiniteAndUnscaleInferMeta
+    param : [x, scale]
+  kernel :
+    func : check_finite_and_unscale
+    param : [x, scale]
+    data_type : x
+  inplace : (x -> out)
+
 - op : cholesky
   args : (Tensor x, bool upper=false)
   output : Tensor

python/paddle/static/amp/amp_nn.py

@@ -41,7 +41,7 @@ def check_finite_and_unscale(x, scale, name=None, float_status=None):
     found_inf = helper.create_variable_for_type_inference(dtype='bool')
 
     if in_dygraph_mode():
-        _C_ops.check_finite_and_unscale_(x, scale, found_inf)
+        x, found_inf = _C_ops.check_finite_and_unscale_(x, scale)
         return x, found_inf
 
     check_type(x, 'x', (tuple, list), 'check_finite_and_unscale')