Method DeviceMemoryFromTensor

tensorflow/compiler/jit/xla_tensor.cc:34–44

static*/

}

/*static*/ se::DeviceMemoryBase XlaTensor::DeviceMemoryFromTensor(
    const Tensor& tensor) {
  const XlaTensor* xla_tensor = FromTensor(&tensor);
  if (xla_tensor) {
    CHECK(xla_tensor->has_shaped_buffer());
    return xla_tensor->shaped_buffer().root_buffer();
  } else {
    return se::DeviceMemoryBase(const_cast<char*>(tensor.tensor_data().data()),
                                tensor.tensor_data().size());
  }
}

Status XlaTensor::AllocateShapedBuffer(DataType dtype,
                                       const xla::Shape& on_host_shape,