auto compilation_result = std::make_unique<XlaCompilationResult>();

  return CompileFromMlirToXlaHlo(
      compile_to_xla_hlo, mlir_to_hlo_args, metadata_proto,
      /*device_type=*/"XLA_TPU_JIT",
      /*shape_determination_fns=*/{}, use_tuple_args, compilation_result.get(),
      custom_legalization_passes, arg_shapes, &arg_core_mapping,
      &per_core_arg_shapes);
}

//   . Legalizes the operations to XLA HLO operations.
//   . Canonicalizes the XLA HLO operations.
//
// device_type: XLA JIT device to use for compilation such as "XLA_CPU_JIT",
//   "XLA_GPU_JIT" or "XLA_TPU_JIT".
// use_tuple_args: when this is true, always create a tuple argument for the
//   entry computation.
// enable_op_fallback: when this is true, prefer tf2xla fallback kernels over
// MLIR

  auto compilation_result = std::make_unique<XlaCompilationResult>();

  return LegalizeTfToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
                         /*device_type=*/"XLA_TPU_JIT",
                         /*shape_determination_fns=*/{}, arg_shapes,
                         &arg_core_mapping, &per_core_arg_shapes,
                         custom_legalization_passes, client,

                   .ok());

  EXPECT_EQ(
      compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
                               mlir::TF::kMlirPh1BridgeCounterV2, "XLA_TPU_JIT",
                               "fallback_disabled", "failure"),
      1);
}

TEST_F(LowerClusterToRuntimeOpsTest, DumpsPipelinePasses) {
  std::vector<std::string> files;

  std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;

  return LegalizeMlirToHlo(mlir_to_hlo_args, metadata_proto, use_tuple_args,
                           /*device_type=*/"XLA_TPU_JIT",
                           custom_legalization_passes,
                           /*shape_determination_fns=*/{}, arg_shapes,
                           &arg_core_mapping, &per_core_arg_shapes, client);
}

        ":xla_cpu_jit",
        "//tensorflow/compiler/plugin",
    ] + if_cuda_or_rocm([
        ":xla_gpu_device",
        ":xla_gpu_jit",
    ]) + if_with_tpu_support([
        ":xla_tpu_device",
        ":xla_tpu_jit",
    ]),
    alwayslink = 1,
)

cc_library(
    name = "xla_cpu_jit",
    visibility = ["//visibility:public"],
    deps = [
        ":jit_compilation_passes",