From f41d1939cfbe2561a6a3f2e10fe0d2be2ed930ac Mon Sep 17 00:00:00 2001 From: Tres Popp Date: Thu, 24 Oct 2019 11:31:54 -0700 Subject: [PATCH] Mark variables initialized in XLA, so Tensorflow also knows they're initialized. PiperOrigin-RevId: 276530886 Change-Id: I5c0d82dce44c81139efc766807fba1039f085de3 --- tensorflow/compiler/jit/xla_launch_util.cc | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/compiler/jit/xla_launch_util.cc b/tensorflow/compiler/jit/xla_launch_util.cc index 76af21b100a..ddaaefcef7d 100644 --- a/tensorflow/compiler/jit/xla_launch_util.cc +++ b/tensorflow/compiler/jit/xla_launch_util.cc @@ -32,6 +32,7 @@ limitations under the License. #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/lib/core/errors.h" @@ -367,6 +368,7 @@ static Status SetBufferForResourceVarTensorUnderAllocateXlaTensors( } } *variable_infos[i].var()->tensor() = output_tensor; + variable_infos[i].var()->is_initialized |= write.modified; return Status::OK(); } @@ -540,6 +542,7 @@ Status XlaComputationLaunchContext::PopulateOutputs( kernel->input_mapping, resource_var_snapshots, write.type, write.shape, buffer, allocator); *variable_infos[i].var()->tensor() = output_tensor; + variable_infos[i].var()->is_initialized |= write.modified; } ++output_num; }