function_(function),
      platform_info_(XlaPlatformInfoFromDevice(ctx->device())),
      has_ref_vars_(has_ref_vars) {}

void XlaLocalLaunchBase::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
  VLOG(1) << "XlaLocalLaunchOpBase::Compute "
          << Canonicalize(function_.name(), AttrSlice(&function_.attr()));
  xla_launch_counter->GetCell(platform_info_.device_type().type_string())
      ->IncrementBy(1);

      GetAllocator(ctx->device(), stream, platform_info_);
  se::DeviceMemoryAllocator* allocator = allocator_ptr.get();
  XlaComputationLaunchContext launch_context(
      client, allocator, client->default_device_ordinal(),
      /*allocate_xla_tensors=*/platform_info_.xla_device_metadata() != nullptr,
      platform_info_.xla_device_metadata()
          ? platform_info_.xla_device_metadata()->UseMultipleStreams()

Status GetCompilationDeviceTypeAndPjRtClient(
    const XlaPlatformInfo& platform_info, FunctionLibraryRuntime* flr,
    DeviceType* compilation_device_type, xla::PjRtClient** pjrt_client) {
  DeviceType device_type = platform_info.device_type();

  if (platform_info.xla_device_metadata()) {
    VLOG(2) << "Building PjRtDeviceCompiler using "
               "platform_info.xla_device_metadata().";

    *compilation_device_type =

  XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(device);

  TF_ASSERT_OK_AND_ASSIGN(
      DeviceType compilation_device_type,
      GetCompilationDeviceType(platform_info.device_type()));

  XlaDeviceCompiler* xla_device_compiler = nullptr;
  TF_EXPECT_OK(BuildXlaDeviceCompiler(device, device_setup_.flr(),
                                      platform_info, compilation_device_type,

  std::shared_ptr<se::DeviceMemoryAllocator> custom_allocator;
  XlaPlatformInfo platform_info(
      compilation_device_type, platform_id, xla_device_metadata.get(),
      /*pjrt_device_metadata=*/nullptr, custom_allocator);

  XlaCompiler::Options options = GenerateCompilerOptionsForPjRt(
      *device_setup_.flr(), device, platform_info,
      /*pjrt_device_compiler=*/nullptr);

                             compiler_arg_source));

  XlaPlatformInfo platform_info = XlaPlatformInfoFromDevice(dev);
  auto compilation_device_type = platform_info.device_type();
  if (platform_info.device_type() != DEVICE_TPU) {
    TF_ASSIGN_OR_RETURN(compilation_device_type,
                        GetCompilationDeviceType(platform_info.device_type()));
  }

  XlaDeviceCompiler* xla_device_compiler;