Prevent copying uninitialized data in AssignOp.

mihaimaruseac · tensorflower-gardener · commit ef1d027be116 · 2021-11-09T11:22:09.000-08:00
This prevents harder to debug undefined behaviors that cannot be traced back to the original tensor after assignments occur earlier in the graph execution. Several of these undefined behaviors are just reference bindings to null pointers, which are caught when running under ubsan/asan.

PiperOrigin-RevId: 408654780
Change-Id: Iad2ec40d43f5fd7ea016c20283356c12d5ddeab1
diff --git a/tensorflow/core/kernels/assign_op.h b/tensorflow/core/kernels/assign_op.h
@@ -50,6 +50,12 @@ class AssignOp : public OpKernel {
     // We always return the input ref.
     context->forward_ref_input_to_ref_output(0, 0);
 
+    // Prevent copying uninitialized data, to solve harder to debug undefined
+    // behaviors that cannot be traced back to the original tensor.
+    OP_REQUIRES(
+        context, rhs.IsInitialized(),
+        errors::Internal("Right hand side of AssignOp is not initialized"));
+
     // We can't always know how this value will be used downstream, so make
     // conservative assumptions in specifying constraints on the memory
     // allocation attributes, unless the Grappler graph analysis determined that
