diff --git a/tensorflow/python/framework/python_op_gen_internal.cc b/tensorflow/python/framework/python_op_gen_internal.cc index eddcd6ac091..c2a47877e31 100644 --- a/tensorflow/python/framework/python_op_gen_internal.cc +++ b/tensorflow/python/framework/python_op_gen_internal.cc @@ -766,36 +766,28 @@ void GenPythonOp::AddDocStringNameArg() { } void GenPythonOp::AddOutputGlobals() { - // Prepare a NamedTuple type to hold the outputs, if there are multiple + // Generate a namedtuple class to hold the outputs, if there are multiple. + // Example: + // + // _OpOutputs = collections.namedtuple( + // "_OpOutputs", + // "out1 out2 out3") if (num_outs_ > 1) { - // Prepare the list of output names - std::vector<string> out_names(num_outs_); + std::vector<string> out_names; + out_names.reserve(num_outs_); for (int i = 0; i < num_outs_; ++i) { - if (!api_def_.out_arg(i).rename_to().empty()) { - out_names[i] = api_def_.out_arg(i).rename_to(); - } else { - out_names[i] = strings::StrCat("output", i); - } + const string out_name = !api_def_.out_arg(i).rename_to().empty() + ? api_def_.out_arg(i).rename_to() + : strings::StrCat("output", i); + out_names.push_back(strings::StrCat("\"", out_name, "\"")); } - string out_names_list = - strings::StrCat("[\"", absl::StrJoin(out_names, "\", \""), "\"]"); - - // Provide the output names as a Python list - string lower_op_name_outputs = - strings::StrCat("_", function_name_, "_outputs"); - const string outputs_prefix = strings::StrCat(lower_op_name_outputs, " = "); - strings::StrAppend(&prelude_, "\n", - WordWrap(outputs_prefix, out_names_list, kRightMargin), - "\n"); strings::StrAppend(&prelude_, "_", op_def_.name(), "Output = collections.namedtuple(\n"); - const string tuple_type_prefix = " "; - const string tuple_type_suffix = strings::StrCat( - "\"", op_def_.name(), "\", ", lower_op_name_outputs, ")"); - strings::StrAppend( - &prelude_, WordWrap(tuple_type_prefix, tuple_type_suffix, kRightMargin), - "\n\n"); + strings::StrAppend(&prelude_, " \"", op_def_.name(), "\",\n"); + strings::StrAppend(&prelude_, " [", absl::StrJoin(out_names, ", "), + "])"); + strings::StrAppend(&prelude_, "\n\n"); } strings::StrAppend(&prelude_, "\n"); }