std::vector<tensorflow::TensorShape> arg_shapes;
  XlaCompilationResult compilation_result;
  auto status = CompileSerializedMlirToXlaHlo(
      kMlirModuleStr, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
      /*use_tuple_args=*/true, /*enable_op_fallback=*/false,
      /*shape_determination_fns=*/{}, &compilation_result);

  EXPECT_TRUE(status.ok());
  EXPECT_THAT(status.value(), HasSubstr("mhlo.const"));
}

syntax = "proto2";

package tensorflow.tf2xla.v2;

// The requested device type to compile for.
enum DeviceType {
  DEVICE_TYPE_UNSPECIFIED = 0;
  XLA_TPU_JIT = 1;
  XLA_CPU_JIT = 2;
  XLA_GPU_JIT = 3;

  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);
}

  auto compilation_result = tensorflow::XlaCompilationResult();

  TF_EXPECT_OK(tensorflow::CompileSerializedMlirToXlaHlo(
                   legalization, arg_shapes, /*device_type=*/"XLA_TPU_JIT",
                   /*use_tuple_args=*/true, /*enable_op_fallback=*/false,
                   /*shape_determination_fns=*/{}, &compilation_result)
                   .status());

//   . 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;

    TF_EXPECT_OK(tensorflow::tf2xla::v2::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)
                     .status());
  }
};

  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",