Function PrepareToUpdateVariable

tensorflow/core/kernels/training_op_helpers.h:198–229

// REQUIRES: If you pass in variable->tensor(), *variable->mu() must be held.
template <typename Device, typename T>
Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor,
                               bool copy_on_read_mode) {
  if (copy_on_read_mode || !tensor->RefCountIsOne()) {
    // Tensor's buffer is in use by some read, so we need to copy before
    // updating.
    PersistentTensor unused;
    Tensor* tmp;
    if (std::is_same<T, Variant>::value) {
      AllocatorAttributes attr;
      attr.set_on_host(true);
      TF_RETURN_IF_ERROR(ctx->allocate_persistent(
          tensor->dtype(), tensor->shape(), &unused, &tmp, attr));

      const auto elements_in = tensor->flat<Variant>();
      auto elements_out = tmp->flat<Variant>();
      for (int64 i = 0; i < elements_in.size(); ++i) {
        elements_out(i) = elements_in(i);
      }
    } else {
      AllocatorAttributes attr;
      attr.set_gpu_compatible(true);
      attr.set_nic_compatible(true);
      TF_RETURN_IF_ERROR(ctx->allocate_persistent(
          tensor->dtype(), tensor->shape(), &unused, &tmp, attr));
      functor::DenseUpdate<Device, T, ASSIGN> copy_functor;
      copy_functor(ctx->eigen_device<Device>(), tmp->flat<T>(),
                   const_cast<const Tensor*>(tensor)->flat<T>());
    }
    *tensor = *tmp;
  }
  return Status::OK();
}

// This gives you `*out`, a tensor you can update, corresponding to a variable
// passed as input index `input`.  This handles the differences between