}

TEST(XlaCompileUtilTest, PjRtXlaLaunchFlagTest) {
  EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

  // Flag is turned on, but no device is allowlisted.
  auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
  rollout_config.enabled_for_xla_launch_ = true;

  EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

    absl::Span<const DataType> result_types);

// Checks if single device compilation and execution with PJRT is enabled for
// `device_type` in either the XlaLaunch op or the XlaCompileOnDemand op.
bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type);

// Gets the resource name of the PjRt DeviceCompiler for `device_type`.
std::string GetPjRtDeviceCompilerResourceName(const DeviceType& device_type);

                        .Attr("index", i)
                        .Finalize(graph.get(), &node);
    TF_RETURN_IF_ERROR(status);
  }
  FixupSourceAndSinkEdges(graph.get());
  return graph;
}

bool UsePjRtForSingleDeviceCompilation(const DeviceType& device_type) {
  const auto& rollout_config = GetXlaOpsCommonFlags()->tf_xla_use_device_api;
  return rollout_config.IsEnabledInXlaLaunchForDevice(device_type) ||

}

Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) {
  if (attr.on_host()) {
    return cpu_allocator();
  }

  if (xla_allocator_ == nullptr) {
    if (UsePjRtForSingleDeviceCompilation(device_name_)) {
      VLOG(1) << "XlaDevice " << this << " uses AsyncValueAllocator";
      pjrt_allocator_ = std::make_unique<AsyncValueAllocator>();
      xla_allocator_ = pjrt_allocator_.get();
    } else {