Fix bugs in AssignOp:

 1. Releasing the unique_ptr would "leak" a TensorBuffer refcount.
 2. The output shape is defined by rhs, not lhs.

PiperOrigin-RevId: 194822802

diff --git a/tensorflow/core/kernels/assign_op.h b/tensorflow/core/kernels/assign_op.h
index 2ed1628..19b38f9 100644 (file)
--- a/tensorflow/core/kernels/assign_op.h
+++ b/tensorflow/core/kernels/assign_op.h
@@ -78,11 +78,10 @@ class AssignOp : public OpKernel {
       // 1. Try to reuse the rhs.
       std::unique_ptr<Tensor> input_alias = context->forward_input(
           1, OpKernelContext::Params::kNoReservation /*output_index*/,
-          old_lhs.dtype(), old_lhs.shape(), DEVICE_MEMORY, attr);
+          rhs.dtype(), rhs.shape(), DEVICE_MEMORY, attr);
       if (input_alias != nullptr) {
         // Transfer ownership to the ref.
-        context->replace_ref_input(0, *input_alias.release(),
-                                   /* lock_held */ true);
+        context->replace_ref_input(0, *input_alias, /* lock_held */ true);
         return;
       }