message: Fix bugs in AssignOp:
  1. Releasing the unique_ptr would "leak" a TensorBuffer refcount.
  2. The output shape is defined by rhs, not lhs.
diff: @@ -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;
  }
  