EXPECT_FALSE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

  // Flag is turned on and the requested device is allowlisted.
  rollout_config.AllowForDeviceInXlaLaunch(DeviceType(DEVICE_CPU));
  EXPECT_TRUE(UsePjRtForSingleDeviceCompilation(DeviceType(DEVICE_CPU)));

  // The requested device is allowlisted, but the flag is turned off.
  rollout_config.enabled_for_xla_launch_ = false;

              StatusIs(error::INTERNAL,
                       HasSubstr(absl::StrCat("PjRtClient for ", DEVICE_CPU,
                                              " is not type PjRtCApiClient"))));
}

TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientSuccess) {
  auto status = pjrt::PjrtApi(DEVICE_CPU);
  if (!status.ok()) {
    TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
  }

TEST_F(XlaCompilerOptionsTest, PjRtOptionsPjRtBaseDevice) {
  // Although DEVICE_CPU isn't a PjRtBaseDevice, we use it here just for testing
  // purposes and to keep things simple. Creating a TpuDevice or
  // NextPluggableDevice in the context of this unit test is non-trivial.
  device_setup_.AddDevicesAndSetUp({DEVICE_CPU});
  Device* device = device_setup_.GetDevice(DEVICE_CPU);
  DeviceType compilation_device_type = DeviceType(DEVICE_CPU_XLA_JIT);

  DeviceType device_gpu = DeviceType(DEVICE_GPU);
  EXPECT_EQ(GetPersistentCacheDirectory(device_gpu), "/tmp/xla_cache");
  DeviceType device_cpu = DeviceType(DEVICE_CPU);
  EXPECT_EQ(GetPersistentCacheDirectory(device_cpu), "/tmp/xla_cache");
  DeviceType device_tpu = DeviceType(DEVICE_TPU);
  EXPECT_TRUE(GetPersistentCacheDirectory(device_tpu).empty());
}

  }

 private:
  TestAsyncIdentityKernel(const TestAsyncIdentityKernel&) = delete;
  void operator=(const TestAsyncIdentityKernel&) = delete;
};

REGISTER_KERNEL_BUILDER(Name("TestAsyncIdentity").Device(DEVICE_CPU),
                        TestAsyncIdentityKernel);

}  // namespace tf_mlrt

};

void TestBitcastOp(Tensor* input_tensor, DataType out_type,
                   TensorShape expected_shape, error::Code expected_code) {
  Status status;
  NodeDef def;
  def.set_op("Bitcast");
  def.set_device(DEVICE_CPU);

  AttrValue typeAttr;
  SetAttrValue(input_tensor->dtype(), &typeAttr);

  AttrValue outTypeAttr;
  SetAttrValue(out_type, &outTypeAttr);

  (*def.mutable_attr())["T"] = typeAttr;

                         error::Code expected_code) {
  // Initialize node used to fetch OpKernel
  Status status;
  NodeDef def;
  def.set_op("ScalarSummary");

  def.set_device(DEVICE_CPU);

  AttrValue valuesTypeAttr;
  SetAttrValue(values->dtype(), &valuesTypeAttr);
  (*def.mutable_attr())["T"] = valuesTypeAttr;

  def.add_input(strings::StrCat("input1: ", DataTypeString(tags->dtype())));

  REGISTER_KERNEL_BUILDER(Name("XlaCallModule").Device(DEVICE),                \
                          XlaCompileOnDemandOp);
REGISTER_XLA_OPS_ON_DEVICE(DEVICE_CPU);
REGISTER_XLA_OPS_ON_DEVICE(DEVICE_GPU);

  int i = 0;
  if (ctx->has_input(i) || ctx->has_input(++i)) {
    ctx->set_output(0, ctx->input(i));
  }
}

REGISTER_KERNEL_BUILDER(Name("XlaLaunch").Device(DEVICE_CPU), XlaLocalLaunchOp);

REGISTER_KERNEL_BUILDER(Name("XlaLaunchV2").Device(DEVICE_CPU), XlaLaunchV2Op);

REGISTER_KERNEL_BUILDER(Name("XlaLaunch")
                            .Device(DEVICE_GPU)
                            .HostMemory("constants")

  CHECK(SerializeToTString(s, output_tstring));
}

void RegisterMergeSummaryOpKernel() {
  TF_Status* status = TF_NewStatus();
  {
    auto* builder = TF_NewKernelBuilder(
        "MergeSummary", tensorflow::DEVICE_CPU, &MergeSummaryOp_Create,
        &MergeSummaryOp_Compute, &MergeSummaryOp_Delete);
    TF_RegisterKernelBuilder("MergeSummary", builder, status);
    CHECK_EQ(TF_OK, TF_GetCode(status))