diff --git a/tf_adapter/kernels/non_zero_with_value_ops.cc b/tf_adapter/kernels/non_zero_with_value_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..b19df065120e18c35664d0199f94f44a8c35fb1f --- /dev/null +++ b/tf_adapter/kernels/non_zero_with_value_ops.cc @@ -0,0 +1,40 @@ +/* + * Copyright (c) Huawei Technologies Co., Ltd. 2019-2020. 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 "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/register_types.h" + +namespace tensorflow { +template +class NonZeroWithValueOP : public OpKernel { + public: + explicit NonZeroWithValueOP(OpKernelConstruction*ctx) : OpKernel(ctx) { + LOG(INFO) << "new NonZeroWithValueOP"; + } + ~NonZeroWithValueOP() { + LOG(INFO) << "del NonZeroWithValueOP"; + } + void Compute(OpKernelContext*ctx) override { + LOG(INFO) << "compute in NonZeroWithValueOP"; + } + bool IsExpensive() override { + LOG(INFO) << "in NonZeroWithValue IsExpensive"; + return false; + } +}; + +REGISTER_KERNEL_BUILDER(Name("NonZeroWithValue").Device(DEVICE_CPU), NonZeroWithValueOP); +} //namespace tensorflow \ No newline at end of file diff --git a/tf_adapter/ops/npu_aicore_ops.cc b/tf_adapter/ops/npu_aicore_ops.cc index c29f84c4052691f9bde100ca4b8bd2be42764bf5..e8899082f7bfd59f7d0336a54e8e1da9def9cc40 100644 --- a/tf_adapter/ops/npu_aicore_ops.cc +++ b/tf_adapter/ops/npu_aicore_ops.cc @@ -341,5 +341,28 @@ REGISTER_OP("NonZero") c->set_output(0, c->MakeShape({rank, -1})); return Status::OK(); }); + +REGISTER_OP("NonZeroWithValue") + .Input("x:T") + .Output("value:T") + .Output("index:output_type") + .Output("count:output_type") + .Attr("transpose:bool = false") + .Attr("T:numbertype") + .Attr("output_type:{int32, int64} = DT_INT64") + .SetIsStateful() + .SetShapeFn([](InferenceContext* c) { + auto input_shape = c->input(0); + int64_t dim1 = c->Value(c->Dim(input_shape, 0)); + int64_t dim2 = c->Value(c->Dim(input_shape, 1)); + int64_t value_num = dim1 * dim2; + int64_t index_dim = 2 * dim1 * dim2; + int64_t count_dim = 1; + + c->set_output(0, c->MakeShape({c->MakeDim(value_num)})); + c->set_output(1, c->MakeShape({c->MakeDim(index_dim)})); + c->set_output(2, c->MakeShape({c->MakeDim(count_dim)})); + return Status::OK(); + }); } // namespace } // namespace tensorflow \ No newline at end of file diff --git a/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py b/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py index 2e03f817f09819c0781963dac2bce0bb7876e930..acaf4485536f7c043b909699a0cf631541153aab 100644 --- a/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py +++ b/tf_adapter/python/npu_bridge/estimator/npu_aicore_ops.py @@ -162,4 +162,16 @@ def nonzero(x, transpose=False, output_type=dtypes.int64, name=None): x = ops.convert_to_tensor(x, name="x") result = npu_aicore_ops.non_zero(x, transpose, output_type, name=name) return result + +def nonzerowithvalue(x, transpose=False, output_type=dtypes.int64, name=None): + """ + nonezero op + Return the indices of the elementes that are non-zero. + Return a tuple of arrays,one for each dimension of a ,containing the indices of the non-zero elementes in that dimension. + The values in a are always tested and returned in row-major ,C-style order. + + """ + x = ops.convert_to_tensor(x, name="x") + result = npu_aicore_ops.non_zero_with_value(x, transpose, output_type, name=name) + return result # go/tf-wildcard-import