Function DLDeviceTypeForDevice

tensorflow/compiler/xla/python/dlpack.cc:194–204

}

StatusOr<DLDeviceType> DLDeviceTypeForDevice(const Device& device) {
  const se::Platform* platform =
      device.local_device_state()->executor()->platform();
  if (platform->id() == se::host::kHostPlatformId) {
    return kDLCPU;
  } else if (platform->id() == se::cuda::kCudaPlatformId) {
    return kDLGPU;
  }
  return InvalidArgument("Device %s cannot be used as a DLPack device.",
                         device.DebugString());
}

StatusOr<DLContext> DLContextForDevice(const Device& device) {
  DLContext context;