// enable logging.
constexpr char kBridgeComponent[] = "TFXLABridge";

using tpu::MlirToHloArgs;
using tpu::ShardingAndIndex;

absl::StatusOr<std::string> CompileFromMlirToXlaHlo(
    bool lower_to_xla_hlo, const MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,

// result of running all the MLIR Bridge passes. If compile_to_xla_hlo is true
// then those passes include all the Legalization to XLA HLO which is returned
// in the compilation_result.
absl::StatusOr<std::string> CompileFromMlirToXlaHlo(
    bool lower_to_xla_hlo, const tpu::MlirToHloArgs& computation,
    const tpu::TPUCompileMetadataProto& metadata, llvm::StringRef device_type,

  LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
                          "Combined MLIR Tf2Xla Bridge.";

  absl::StatusOr<std::string> mlir_compilation =
      internal::CompileFromMlirToXlaHlo(
          /*lower_to_xla_hlo=*/false, computation, metadata, device_type,
          shape_determination_fns, use_tuple_args, compilation_result,
          custom_legalization_passes, arg_shapes, arg_core_mapping,

  std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
  std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;

  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(),