From b7b1a1e0c7c261f40a76c2e1c0522b313b3e1509 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B3=8A=E9=9C=86?= <hujunqi.hjq@alibaba-inc.com>
Date: Wed, 10 Jan 2024 14:13:27 +0800
Subject: [PATCH] [Embedding] modify UniqueWithExtraCounts name.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: 泊霆 <hujunqi.hjq@alibaba-inc.com>
---
 tensorflow/core/kernels/unique_ali_op.cc | 8 ++++----
 tensorflow/core/ops/array_ops.cc         | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/unique_ali_op.cc b/tensorflow/core/kernels/unique_ali_op.cc
index 6f8e3d1bb76..57dea40ac96 100644
--- a/tensorflow/core/kernels/unique_ali_op.cc
+++ b/tensorflow/core/kernels/unique_ali_op.cc
@@ -208,12 +208,12 @@ class UniqueWithCountAliOp : public UniqueAliOp<T, TIndex> {
                               .TypeConstraint<type>("T")         \
                               .TypeConstraint<int64>("out_idx"), \
                           UniqueAliOp<type, int64>)              \
-  REGISTER_KERNEL_BUILDER(Name("UniqueWithExtraCounts")          \
+  REGISTER_KERNEL_BUILDER(Name("_UniqueWithExtraCounts")         \
                               .Device(DEVICE_CPU)                \
                               .TypeConstraint<type>("T")         \
                               .TypeConstraint<int32>("out_idx"), \
                           UniqueWithCountAliOp<type, int32>)     \
-  REGISTER_KERNEL_BUILDER(Name("UniqueWithExtraCounts")          \
+  REGISTER_KERNEL_BUILDER(Name("_UniqueWithExtraCounts")         \
                               .Device(DEVICE_CPU)                \
                               .TypeConstraint<type>("T")         \
                               .TypeConstraint<int64>("out_idx"), \
@@ -242,12 +242,12 @@ REGISTER_UNIQUE(string)
                               .TypeConstraint<type>("T")         \
                               .TypeConstraint<int64>("out_idx"), \
                           UniqueAliOp<type, int64>)              \
-  REGISTER_KERNEL_BUILDER(Name("UniqueWithExtraCounts")          \
+  REGISTER_KERNEL_BUILDER(Name("_UniqueWithExtraCounts")         \
                               .Device(DEVICE_GPU)                \
                               .TypeConstraint<type>("T")         \
                               .TypeConstraint<int32>("out_idx"), \
                           UniqueWithCountAliOp<type, int32>)     \
-  REGISTER_KERNEL_BUILDER(Name("UniqueWithExtraCounts")          \
+  REGISTER_KERNEL_BUILDER(Name("_UniqueWithExtraCounts")         \
                               .Device(DEVICE_GPU)                \
                               .TypeConstraint<type>("T")         \
                               .TypeConstraint<int64>("out_idx"), \
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 306026977ef..71295a0ecfd 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -1743,7 +1743,7 @@ REGISTER_OP("UniqueWithCountsV2")
 
 // ---------------------------------------------------
 
-REGISTER_OP("UniqueWithExtraCounts")
+REGISTER_OP("_UniqueWithExtraCounts")
     .Input("x: T")
     .Input("extra_indices: N * T")
     .Input("extra_counts: N * out_idx")