diff --git a/tensorflow/compiler/xla/protobuf_util.cc b/tensorflow/compiler/xla/protobuf_util.cc index e476015f94f..b7c30531923 100644 --- a/tensorflow/compiler/xla/protobuf_util.cc +++ b/tensorflow/compiler/xla/protobuf_util.cc @@ -39,12 +39,17 @@ bool ProtobufEquals(const tensorflow::protobuf::Message& m1, } Status DumpProtoToDirectory(const tensorflow::protobuf::Message& message, - const string& directory, const string& file_name) { + const string& directory, const string& file_name, + string* full_path) { tensorflow::Env* env = tensorflow::Env::Default(); TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(directory)); string safe_file_name = SanitizeFileName(file_name) + ".pb"; - const string path = tensorflow::io::JoinPath(directory, safe_file_name); - return tensorflow::WriteBinaryProto(env, path, message); + string full_path_impl; + if (!full_path) { + full_path = &full_path_impl; + } + *full_path = tensorflow::io::JoinPath(directory, safe_file_name); + return tensorflow::WriteBinaryProto(env, *full_path, message); } } // namespace protobuf_util diff --git a/tensorflow/compiler/xla/protobuf_util.h b/tensorflow/compiler/xla/protobuf_util.h index e20a7e95a63..7db020982b9 100644 --- a/tensorflow/compiler/xla/protobuf_util.h +++ b/tensorflow/compiler/xla/protobuf_util.h @@ -37,8 +37,12 @@ extern bool ProtobufEquals(const tensorflow::protobuf::Message& m1, // 'directory/file_name.pb'. The 'directory' is recursively created if it // doesn't already exist, and the 'file_name' is sanitized by replacing // illegal characters with underscore '_'. +// +// If 'full_name' is not null then it is set to the name of the file the +// protobuf was written to. Status DumpProtoToDirectory(const tensorflow::protobuf::Message& message, - const string& directory, const string& file_name); + const string& directory, const string& file_name, + string* full_path = nullptr); // Registers a function that may either expand a dirpath or forward the original // dirpath along as-is.