return %0 : tensor<1x3xf32>
    }
  )mlir";

  const OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kModuleWithCompositeDotGeneral);
  ASSERT_TRUE(module_op);

  // Create a pass manager with `SaveQuantizationReportInstrumentation` and
  // `QuantizeCompositeFunctionsPass`. Run the passes against `module_op`.
  PassManager pm(ctx_.get());

  QuantizeCompositeFunctionsPassOptions options;

absl::Status RunPasses(const absl::string_view name, FuncT add_passes_func,
                       mlir::MLIRContext& ctx, mlir::ModuleOp module_op) {
  mlir::PassManager pm{&ctx};
  add_passes_func(pm);

  mlir::StatusScopedDiagnosticHandler diagnostic_handler{&ctx};
  TF_RETURN_IF_ERROR(MaybeEnableIrPrinting(pm, name));

  if (failed(pm.run(module_op))) {
    return absl::InternalError(
        absl::StrFormat("Failed to run pass: %s. %s", name,

PostCalibrationComponent::PostCalibrationComponent(
    absl::Nonnull<MLIRContext*> ctx)
    : ctx_(ABSL_DIE_IF_NULL(ctx)) {}  // Crash OK

absl::StatusOr<ModuleOp> PostCalibrationComponent::Run(
    ModuleOp module_op, const QuantizationConfig& config) {
  TF_RETURN_IF_ERROR(RunPasses(
      kName, /*add_passes_func=*/
      [&config](PassManager& pm) {

namespace stablehlo::quantization {

// Disables debugging on `DumpTensor` ops.
void DisableDebugging(mlir::ModuleOp module_op);

// Changes the filename from `unquantized_tensor_data.pb` to
// `quantized_tensor_data.pb`.
void ChangeToQuantizedFilename(mlir::ModuleOp module_op);

}  // namespace stablehlo::quantization

void UpdateFunctionAliases(
    absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
    ModuleOp module_op);

// Loads a SavedModel to `mlir::ModuleOp` and performs preprocesses including
// shape inference and graph freezing.
// TODO: b/329206105 - Add unit tests after decomposing preprocessing passes.
absl::StatusOr<OwningOpRef<ModuleOp>> ImportSavedModel(
    absl::string_view saved_model_path,

absl::Status AddCalibrationStatistics(
    mlir::ModuleOp module_op, absl::string_view calibration_data_dir,
    const CalibrationOptions& calibration_options,
    const PyFunctionLibrary& py_function_library) {
  TF_ASSIGN_OR_RETURN(const CalibrationStatisticsFlatMap statistics_map,
                      ReadStatistics(calibration_data_dir));

  absl::Status status = absl::OkStatus();
  module_op.walk([&py_function_library, &calibration_options, &status,

                                                         Operation* op) {
  // Only run after `QuantizeCompositeFunctionPass`.
  if (!IsQuantizeCompositeFunctionPass(pass, op)) return;

  auto module_op = cast<ModuleOp>(op);
  const QuantizationReport report(module_op);

  // Print a human-readable report to stdout regardless of whether the report
  // is saved to file.
  report.Print();

  // Exit early if the report should not be saved to file.

      /*noinline_functions=*/aliased_function_names, *module_op, context,
      session,
      /*run_tf_to_stablehlo=*/true, /*deserialize_xla_call_module=*/false,
      input_arg_shapes));

  return std::move(module_op);
}

}  // namespace

absl::StatusOr<OwningOpRef<ModuleOp>> TfToStablehlo(
    absl::string_view input_path, MLIRContext* context,

namespace tensorflow {
namespace quantization {

absl::Status RunPassesOnModuleOp(
    std::optional<absl::string_view> mlir_dump_file_name,
    mlir::PassManager& pass_manager, mlir::ModuleOp module_op) {
  mlir::StatusScopedDiagnosticHandler statusHandler(module_op.getContext(),
                                                    /*propagate=*/true);

  absl::StatusOr<std::unique_ptr<llvm::raw_ostream>> dump_file;
  if (mlir_dump_file_name) {

  }
  return success();
}

// Fills `op_to_parallel_ids_map` from parallel execution attributes in
// `module_op`. Returns `failure` iff any attribute is malformed.
LogicalResult FillOpToParallelIdsMap(
    ModuleOp module_op, TF::OpToParallelIdsMap& op_to_parallel_ids_map) {
  auto result = module_op->walk([&](GraphOp graph) {
    if (failed(FillOpToParallelIdsMap(graph, op_to_parallel_ids_map)))