Method MakeTensorFromProto

tensorflow/compiler/jit/xla_device.cc:488–518

}

Status XlaDevice::MakeTensorFromProto(XlaDeviceContext* device_context,
                                      const TensorProto& tensor_proto,
                                      const AllocatorAttributes alloc_attrs,
                                      Tensor* tensor) {
  Tensor parsed(tensor_proto.dtype());
  if (!parsed.FromProto(cpu_allocator(), tensor_proto)) {
    return errors::InvalidArgument("Cannot parse tensor from proto: ",
                                   tensor_proto.DebugString());
  }

  Status status;
  if (alloc_attrs.on_host()) {
    *tensor = parsed;
  } else {
    mutex_lock lock(mu_);
    Allocator* allocator = GetAllocatorLocked(alloc_attrs);
    Tensor copy(allocator, parsed.dtype(), parsed.shape());
    Notification n;
    device_context->CopyCPUTensorToDevice(
        &parsed, this, &copy,
        [&n, &status](const Status& s) {
          status = s;
          n.Notify();
        },
        true /*sync_dst_compute*/);
    n.WaitForNotification();
    *tensor = copy;
  }
  VLOG(2) << "Allocated tensor at " << DMAHelper::base(tensor);
  return status;
}

Status XlaDevice::MakeTensorFromProto(const TensorProto& tensor_proto,
                                      const AllocatorAttributes alloc_attrs,