| 63 | } |
| 64 | |
| 65 | StatusOr<std::unique_ptr<HloModule>> LoadModuleFromData( |
| 66 | const string& data, const string& format, |
| 67 | hlo_module_loader_details::Config ovr_config, |
| 68 | const std::function<void(HloModuleConfig*)>& config_modifier_hook) { |
| 69 | DebugOptions debug_options = GetDebugOptionsFromFlags(); |
| 70 | std::unique_ptr<HloModule> module; |
| 71 | if (format == "hlo" || format == "txt") { |
| 72 | string hlo_string = StripLogHeaders(data); |
| 73 | HloModuleConfig config; |
| 74 | config.set_debug_options(debug_options); |
| 75 | TF_RETURN_IF_ERROR(OverrideConfig(ovr_config, &config)); |
| 76 | if (config_modifier_hook) { |
| 77 | config_modifier_hook(&config); |
| 78 | } |
| 79 | TF_ASSIGN_OR_RETURN(module, |
| 80 | ParseAndReturnUnverifiedModule(hlo_string, config)); |
| 81 | } else { |
| 82 | HloSnapshot proto; |
| 83 | if (format == "pb") { |
| 84 | if (!proto.ParseFromString(data) && |
| 85 | !proto.mutable_hlo()->ParseFromString(data) && |
| 86 | !proto.mutable_hlo()->mutable_hlo_module()->ParseFromString(data)) { |
| 87 | return InvalidArgument("Failed to parse input as HLO protobuf binary"); |
| 88 | } |
| 89 | } else if (format == "pbtxt") { |
| 90 | if (!google::protobuf::TextFormat::ParseFromString(data, &proto) && |
| 91 | !google::protobuf::TextFormat::ParseFromString(data, proto.mutable_hlo()) && |
| 92 | !google::protobuf::TextFormat::ParseFromString( |
| 93 | data, proto.mutable_hlo()->mutable_hlo_module())) { |
| 94 | return InvalidArgument("Failed to parse input as HLO protobuf text"); |
| 95 | } |
| 96 | } else { |
| 97 | return InvalidArgument( |
| 98 | "Invalid format from file extension: '%s'. Expected: hlo, txt, pb, " |
| 99 | "or pbtxt", |
| 100 | format); |
| 101 | } |
| 102 | TF_ASSIGN_OR_RETURN(HloModuleConfig config, |
| 103 | HloModule::CreateModuleConfigFromProto( |
| 104 | proto.hlo().hlo_module(), debug_options)); |
| 105 | TF_RETURN_IF_ERROR(OverrideConfig(ovr_config, &config)); |
| 106 | if (config_modifier_hook) { |
| 107 | config_modifier_hook(&config); |
| 108 | } |
| 109 | TF_ASSIGN_OR_RETURN( |
| 110 | module, HloModule::CreateFromProto(proto.hlo().hlo_module(), config)); |
| 111 | } |
| 112 | return std::move(module); |
| 113 | } |
| 114 | |
| 115 | StatusOr<std::unique_ptr<HloModule>> LoadModuleFromFile( |
| 116 | const string& path, hlo_module_loader_details::Config ovr_config, |