experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Support loading HloModuleProto using xla::LoadModuleFromData

Browse Source 
PiperOrigin-RevId: 291743207
Change-Id: I44fd0e5151075f969b7d333f0c3d1e7dd06e6237
This commit is contained in:
A. Unique TensorFlower 2020-01-27 09:56:28 -08:00 committed by TensorFlower Gardener
parent 5adb929c26
commit f9533a9c49
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/compiler/xla/tools/hlo_module_loader.cc
View File

@ -82,12 +82,15 @@ StatusOr<std::unique_ptr<HloModule>> LoadModuleFromData(
HloSnapshot proto;
if (format == "pb") {
if (!proto.ParseFromString(data) &&
!proto.mutable_hlo()->ParseFromString(data)) {
!proto.mutable_hlo()->ParseFromString(data) &&
!proto.mutable_hlo()->mutable_hlo_module()->ParseFromString(data)) {
return InvalidArgument("Failed to parse input as HLO protobuf binary");
}
} else if (format == "pbtxt") {
if (!google::protobuf::TextFormat::ParseFromString(data, &proto) &&
!google::protobuf::TextFormat::ParseFromString(data, proto.mutable_hlo())) {
!google::protobuf::TextFormat::ParseFromString(data, proto.mutable_hlo()) &&
!google::protobuf::TextFormat::ParseFromString(
data, proto.mutable_hlo()->mutable_hlo_module())) {
return InvalidArgument("Failed to parse input as HLO protobuf text");
}
} else {