// Overload of `PreprocessAndFreezeGraph` that uses the default MLIR dump file
// prefix.
inline absl::Status PreprocessAndFreezeGraph(mlir::ModuleOp module_op,
                                             mlir::MLIRContext* context,
                                             std::optional<Session*> session) {
  return PreprocessAndFreezeGraph(
      /*mlir_dump_file_prefix=*/kDefaultTfQuantMlirDumpFilePrefix,

using ::stablehlo::quantization::QuantizationConfig;
using ::tensorflow::MLIRImportOptions;
using ::tensorflow::SavedModelBundle;
using ::tensorflow::SavedModelSignatureDefsToMlirImport;
using ::tensorflow::quantization::PreprocessAndFreezeGraph;

absl::StatusOr<ImportedMlirModuleOp> SavedModelToMlirModuleOp(
    const absl::string_view saved_model_path,
    const std::unordered_set<std::string>& tags,
    const std::vector<std::string>& signature_keys,

  llvm::SmallVector<llvm::ArrayRef<int64_t>> input_arg_shapes(
      input_arg_shapes_vec.begin(), input_arg_shapes_vec.end());
  TF_RETURN_IF_ERROR(tensorflow::quantization::PreprocessAndFreezeGraph(
      /*mlir_dump_file_prefix=*/"", /*is_inliner_run=*/true,
      /*noinline_functions=*/aliased_function_names, *module_op, context,
      session,

  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
  // MHLO -> StableHLO legalization.
  pm.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());
}

absl::Status PreprocessAndFreezeGraph(
    const absl::string_view mlir_dump_file_prefix, const bool is_inliner_run,
    const absl::flat_hash_set<std::string>& noinline_functions,
    mlir::ModuleOp module_op, mlir::MLIRContext* context,

      LOG(INFO) << "file type is mlir. ignoring freezing TF graph";
    } else {
      LOG(INFO) << "--freeze-tf-graph is on. Freezing TF graph...";
      const auto status = tensorflow::quantization::PreprocessAndFreezeGraph(
          *module, module->getContext(), session);
      if (!status.ok()) {
        return tensorflow::errors::Aborted(
            "Failed to preprocess & freeze TF graph");
      }
    }
  }

  absl::flat_hash_set<std::string> aliased_function_names;
  absl::c_for_each(function_aliases, [&](const auto &aliases) {
    return aliased_function_names.insert(aliases.first);
  });

  TF_RETURN_IF_ERROR(PreprocessAndFreezeGraph(
      /*mlir_dump_file_prefix=*/kDefaultTfQuantMlirDumpFilePrefix,
      /*is_inliner_run=*/is_inliner_run,
      /*noinline_functions=*/aliased_function_names, module_ref.get(), context,

  // stablehlo, we automatically enable dynamic range quantization

  if (toco_flags.post_training_quantize()) {
    const auto status = quantization::PreprocessAndFreezeGraph(
        module, module.getContext(), session);
    if (!status.ok()) {
      return status_handler.Combine(
          absl::InternalError("Failed to preprocess & freeze TF graph."));
    }