llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes,
    llvm::StringRef module_name = llvm::StringRef(),
    bool lower_to_xla_hlo = true);

// Runs MLIR Bridge on a serialized MLIR module.
//
// If lower_to_xla_hlo is true then compiles down into XLA HLO, generates all
// accompanying metadata and stores them in CompilationResult.

void AddLegalizationPasses(mlir::OpPassManager& pm, bool legalize_chlo,
                           llvm::StringRef device_type, bool enable_op_fallback,
                           bool lower_to_xla_hlo) {
  if (lower_to_xla_hlo) {
    // Lower TF quant ops and types to MHLO int.
    mlir::quant::stablehlo::AddQuantizationLoweringPasses(pm);

    pm.addPass(mlir::mhlo::createLegalizeTFPass(
        legalize_chlo,

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

                                    /*enable_op_fallback=*/false,
                                    /*custom_legalization_passes*/ {},
                                    /*lower_to_xla_hlo=*/false);

  std::string pass_description;
  llvm::raw_string_ostream raw_stream(pass_description);
  pass_manager.printAsTextualPipeline(raw_stream);

// 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,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,

                          "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,
          per_core_arg_shapes);