commit f9533a9c49b6b183a6a562f41e3fdeccfc54d89b
author A. Unique TensorFlower <gardener@tensorflow.org> Mon Jan 27 09:56:28 2020 -0800
committer TensorFlower Gardener <gardener@tensorflow.org> Mon Jan 27 09:59:58 2020 -0800
tree 6014e53483c2049f17d3ad633ced11283fac188e
parent 5adb929c26ce732802ec76d2a094e479543da31b

Support loading HloModuleProto using xla::LoadModuleFromData

PiperOrigin-RevId: 291743207
Change-Id: I44fd0e5151075f969b7d333f0c3d1e7dd06e6237

tensorflow/compiler/xla/tools/hlo_module_loader.cc

diff --git a/tensorflow/compiler/xla/tools/hlo_module_loader.cc b/tensorflow/compiler/xla/tools/hlo_module_loader.cc
index 0b16c87..b3aaba7 100644
--- a/tensorflow/compiler/xla/tools/hlo_module_loader.cc
+++ b/tensorflow/compiler/xla/tools/hlo_module_loader.cc
@@ -82,12 +82,15 @@
     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 {