[Embedding] modify UniqueWithExtraCounts name.

Signed-off-by: 泊霆 <[email protected]>
DeepRec-AI · Jan 10, 2024 · b7b1a1e · b7b1a1e
1 parent b75c8d8
commit b7b1a1e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git 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
@@ -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")