diff --git a/tensorflow/core/profiler/utils/tf_op_utils.cc b/tensorflow/core/profiler/utils/tf_op_utils.cc index 43503906f67..525d18b68c6 100644 --- a/tensorflow/core/profiler/utils/tf_op_utils.cc +++ b/tensorflow/core/profiler/utils/tf_op_utils.cc @@ -42,7 +42,8 @@ const absl::string_view kMemcpyHToDOp = "MemcpyHToD"; const absl::string_view kMemcpyDToHOp = "MemcpyDToH"; bool IsTfOpName(absl::string_view op_name) { - static const LazyRE2 kTfOpNameRegEx = {"[A-Za-z0-9.][A-Za-z0-9_./]*"}; + // TODO(b/177602927): Confirm the naming convention with the TF team. + static const LazyRE2 kTfOpNameRegEx = {"[A-Za-z0-9.][A-Za-z0-9_.\\/>-]*"}; return RE2::FullMatch(op_name, *kTfOpNameRegEx); } diff --git a/tensorflow/core/profiler/utils/tf_op_utils_test.cc b/tensorflow/core/profiler/utils/tf_op_utils_test.cc index edefe81c9cd..d0dce8fe3dc 100644 --- a/tensorflow/core/profiler/utils/tf_op_utils_test.cc +++ b/tensorflow/core/profiler/utils/tf_op_utils_test.cc @@ -157,6 +157,18 @@ TEST(TfOpUtilsTest, OpWithoutTypeTest) { EXPECT_EQ(TfOpEventName(kName), "OpName"); // without trailing ':' } +TEST(TfOpUtilsTest, NameScopeTest) { + const absl::string_view kName = "scope-1/scope2/OpName:OpType"; + TfOp tf_op = ParseTfOpFullname(kName); + EXPECT_EQ(tf_op.category, Category::kTensorFlow); + EXPECT_EQ(tf_op.name, "scope-1/scope2/OpName"); + EXPECT_EQ(tf_op.type, "OpType"); + std::vector<absl::string_view> name_scopes = ParseTfNameScopes(tf_op); + EXPECT_EQ(name_scopes.size(), 2); + EXPECT_EQ(name_scopes[0], "scope-1"); + EXPECT_EQ(name_scopes[1], "scope2"); +} + } // namespace } // namespace profiler } // namespace tensorflow