}
  }

  return WalkResult::advance();
}

FailureOr<OwningOpRef<ModuleOp>> CreatePrunedModule(
    ModuleOp module, llvm::ArrayRef<llvm::StringRef> function_names) {
  SymbolTableCollection symbol_table;
  OpBuilder builder(module.getContext());

  OwningOpRef<ModuleOp> pruned =
      builder.create<ModuleOp>(module->getLoc());
  (*pruned)->setAttrs(module->getAttrs());

  // Parses `module_op_str` to create a `ModuleOp`.
  OwningOpRef<ModuleOp> ParseModuleOpString(
      const absl::string_view module_op_str) {
    return parseSourceString<ModuleOp>(module_op_str, ctx_.get());
  }

  // Convenience function that returns the first operation of type `OpT` from
  // the `@main` function in `module_op`. Useful when testing with a text
  // representation of a `ModuleOp` containing a single function `@main`.

  }];
}

def RestoreFunctionNamePass : Pass<"stablehlo-restore-function-name", "ModuleOp"> {
  let summary = "Restores function name from XlaCallModule op.";
}

def QuantizeCompositeFunctionsPass : Pass<"stablehlo-quantize-composite-functions", "ModuleOp"> {
  let summary = "Quantize composite functions with QDQ input / outputs.";
  let options = [

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

TEST_F(IsOpQuantizableStableHloTest, ConstantOpQuantizable) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
  ASSERT_TRUE(module_op);

  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
  ASSERT_THAT(test_func, NotNull());

  auto constant_op =

namespace tensorflow {
namespace {

absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GetMlirModuleFromString(
    llvm::StringRef string, mlir::MLIRContext* context) {
  mlir::DialectRegistry mlir_registry;
  RegisterAllTensorFlowDialects(mlir_registry);
  context->appendDialectRegistry(mlir_registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
  auto status =

// Converts a TensorFlow GraphDef contained in `input` param into a MLIR module.
// Creates MLIR entities into the given MLIR `context`.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
GraphdefToMlirTranslateFunction(
    llvm::StringRef input, const std::vector<std::string>& input_arrays,
    const std::vector<std::string>& input_dtypes,
    const std::vector<std::optional<std::vector<int>>>& input_shapes,

absl::Status CalibrationComponent::ExportToSavedModel(
    ModuleOp module_op, absl::string_view calibration_data_dir,
    const bool force_regenerate_calibration_data,
    const absl::string_view dst_saved_model_path) {
  TF_ASSIGN_OR_RETURN(const std::string checkpoint_dir, GetLocalTmpFileName());

  // Clone ModuleOp and function aliases so changes in this pipeline won't
  // be reflected in the original values.

                                     absl::Nullable<Operation*> op) {
  // It is known that `op` is `ModuleOp` when `pass` is
  // `QuantizeCompositeFunctionPass`, but the check is still performed to be
  // defensive.
  return pass != nullptr &&
         pass->getArgument() == "stablehlo-quantize-composite-functions" &&
         isa_and_nonnull<ModuleOp>(op);
}

// Report is saved only when:
// * After running `QuantizeCompositeFunctionPass`.

#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/platform/errors.h"

namespace tensorflow {

std::string SerializeMlirModule(mlir::ModuleOp module_op) {
  std::string serialized_mlir_module;
  llvm::raw_string_ostream os(serialized_mlir_module);
  mlir::OpPrintingFlags print_flags;
  print_flags.enableDebugInfo();
  module_op.print(os, print_flags);

    return tmp_dir;
  }

  // Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
  // module op is valid.
  mlir::OwningOpRef<mlir::ModuleOp> ParseModuleOpString(
      const absl::string_view module_op_str) {
    auto module_op_ref =
        mlir::parseSourceString<mlir::ModuleOp>(module_op_str, &ctx_);
    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }