std::unique_ptr<Tensor> input_alias =
context->forward_input(1, dtype_, value.shape(), DEVICE_MEMORY, attr);
mutex_lock ml(*variable->mu());
+ variable->is_initialized = true;
if (input_alias) {
*variable->tensor() = *input_alias;
return;
DataTypeString(DT_VARIANT)));
mutex_lock ml(*variable->mu());
-
+ variable->is_initialized = true;
*variable->tensor() = Tensor(DT_VARIANT, value.shape());
const auto elements_in = value.flat<Variant>();
auto elements_out = variable->tensor()->flat<Variant>();
#undef REGISTER_GPU_KERNELS
#endif // GOOGLE_CUDA
+class VarIsInitializedOp : public OpKernel {
+ public:
+ explicit VarIsInitializedOp(OpKernelConstruction* c) : OpKernel(c) {}
+
+ void Compute(OpKernelContext* context) override {
+ Tensor* output = nullptr;
+ OP_REQUIRES_OK(context,
+ context->allocate_output(0, TensorShape({}), &output));
+ auto output_tensor = output->tensor<bool, 0>();
+ Var* variable = nullptr;
+ Status s = LookupResource(context, HandleFromInput(context, 0), &variable);
+ if (!s.ok()) {
+ output_tensor() = false;
+ return;
+ }
+ core::ScopedUnref su(variable);
+ mutex_lock ml(*variable->mu());
+ output_tensor() = variable->is_initialized;
+ }
+};
+
REGISTER_KERNEL_BUILDER(Name("VarIsInitializedOp").Device(DEVICE_CPU),
- IsResourceInitialized<Var>);
+ VarIsInitializedOp);
#if GOOGLE_CUDA
REGISTER_KERNEL_BUILDER(Name("VarIsInitializedOp")
tensor_.shape().DebugString());
}
+ // Only used in the resource variable path. In resource variables,
+ // tensor.IsInitialized() can be true (i.e. have memory allocated to it) while
+ // there is not a good value there due to a race condition, and it's possible
+ // to stumble upon this during variable.initialized_value(). So it's best to
+ // just store directly whether the variable is initialized.
+ bool is_initialized = false; // GUARDED_BY(mu_) but annotalysis doesn't like
+ // it.
+
private:
mutex mu_;
Tensor tensor_;