diff options
Diffstat (limited to 'tensorflow/compiler/aot/compile.cc')
-rw-r--r-- | tensorflow/compiler/aot/compile.cc | 10 |
1 files changed, 7 insertions, 3 deletions
diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index eac8da0ab1..2b8cc6024c 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -97,11 +97,15 @@ Status CompileGraph(const GraphDef& graph_def, const tf2xla::Config& config, TF_RETURN_IF_ERROR(ConvertGraphDefToXla(graph_def, config, client, &computation, &compile_result->has_context_arg)); - if (!flags.debug_dir.empty()) { + if (!flags.out_session_module.empty()) { TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::SessionModule> module, computation.Snapshot()); - string file = io::JoinPath(flags.debug_dir, "tfcompile_xla_module.pb"); - TF_RETURN_IF_ERROR(WriteBinaryProto(Env::Default(), file, *module)); + // Serialize the SessionModule deterministically so that all the outputs of + // a tf_library genrule are deterministic. + string proto; + TF_RET_CHECK(SerializeToStringDeterministic(*module, &proto)); + TF_RETURN_IF_ERROR( + WriteStringToFile(Env::Default(), flags.out_session_module, proto)); } xla::cpu::CpuAotCompilationOptions aot_opts( flags.target_triple, flags.target_cpu, flags.target_features, |