support auto generate for nonzero (#51600)

paddle/fluid/operators/where_index_op.cc

-/* Copyright (c) 2019 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/unary.h"
-
-namespace paddle {
-namespace operators {
-
-class WhereIndexOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "Condition");
-    return phi::KernelKey(data_type, ctx.GetPlace());
-  }
-};
-
-class WhereIndexOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() override {
-    AddInput("Condition", "A bool tensor whose rank is at least 1");
-    AddOutput("Out", "An int64 tensor of rank 2");
-    AddComment(R"DOC(
-      Return a int64 tensor with rank 2, specifying the coordinate of true element in `Condition`.
-)DOC");
-  }
-};
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-DECLARE_INFER_SHAPE_FUNCTOR(where_index,
-                            WhereIndexInferShapeFunctor,
-                            PD_INFER_META(phi::NonZeroInferMeta));
-REGISTER_OPERATOR(
-    where_index,
-    ops::WhereIndexOp,
-    ops::WhereIndexOpMaker,
-    paddle::framework::EmptyGradOpMaker<paddle::framework::OpDesc>,
-    paddle::framework::EmptyGradOpMaker<paddle::imperative::OpBase>,
-    WhereIndexInferShapeFunctor);

paddle/phi/api/yaml/legacy_ops.yaml

@@ -1208,14 +1208,6 @@
     func : nms
     data_type : x
 
-- op : nonzero
-  args : (Tensor condition)
-  output : Tensor(out)
-  infer_meta :
-    func : NonZeroInferMeta
-  kernel :
-    func : nonzero
-
 - op : norm
   args : (Tensor x, int axis, float epsilon, bool is_test)
   output : Tensor(out), Tensor(norm)

paddle/phi/api/yaml/ops.yaml

@@ -1011,6 +1011,15 @@
   optional : weight
   backward : nll_loss_grad
 
+- op : nonzero
+  args : (Tensor condition)
+  output : Tensor(out)
+  infer_meta :
+    func : NonZeroInferMeta
+  kernel :
+    func : nonzero
+    data_type: condition
+
 - op : npu_identity
   args : (Tensor x, int format = -1)
   output : Tensor

paddle/phi/ops/compat/where_index_sig.cc

-// Copyright (c) 2022 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/phi/core/compat/op_utils.h"
-
-namespace phi {
-
-KernelSignature WhereIndexOpArgumentMapping(const ArgumentMappingContext& ctx) {
-  return KernelSignature("nonzero", {"Condition"}, {}, {"Out"});
-}
-
-}  // namespace phi
-
-PD_REGISTER_BASE_KERNEL_NAME(where_index, nonzero);
-
-PD_REGISTER_ARG_MAPPING_FN(where_index, phi::WhereIndexOpArgumentMapping);