// passed to the target function.
LogicalResult LiftHashTableOpsToArguments(ModuleOp module_op,
                                          func::FuncOp target_func) {
  if (!llvm::hasSingleElement(target_func)) return success();
  if (!UsedBySupportedOps(module_op, target_func)) return success();
  if (IsMainOrInitializerFunction(module_op, target_func)) return success();

  llvm::StringMap<int> shared_name_to_arg_idx;

}  // namespace

absl::Status RunCalibrationPasses(
    mlir::ModuleOp module_op, MLIRContext& ctx,
    absl::string_view calibration_data_dir,
    const bool force_regenerate_calibration_data) {
  // Disable DumpTensor ops when running calibration.
  DisableDebugging(module_op);

  std::vector<std::string> skipping_aggregator_ops;
  if (!force_regenerate_calibration_data) {

  TF_RETURN_IF_ERROR(RefineShapes(arg_shapes, module_op));

  if (VLOG_IS_ON(2))
    tensorflow::DumpMlirOpToFile("compile_mlir_shape_refiner", module_op);

  return absl::OkStatus();
}

Status BuildHloFromTf(mlir::ModuleOp module_op, xla::XlaBuilder& builder,
                      llvm::ArrayRef<xla::XlaOp> xla_params,
                      std::vector<xla::XlaOp>& returns,

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

TEST_F(LiftAsFunctionCallTest, LiftedFunctionSucceeds) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleLifted);
  ASSERT_TRUE(module_op);

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

SessionInitializerOp GetSessionInitializerOp(ModuleOp module_op);

// Returns the exported name for the session initializer function.
SmallVector<StringRef, 2> GetSessionInitializerExportedName(ModuleOp module_op);

// Returns initializer function ops. These functions' symbols are in the
// "initializers" attribute of the session initializer op.
SmallVector<func::FuncOp, 2> GetInitializerFunctions(ModuleOp module_op);

SessionInitializerOp GetOrCreateSessionInitializerOp(ModuleOp module_op) {
  SessionInitializerOp session_init_op = GetSessionInitializerOp(module_op);

  // Create one if it doesn't exist.
  if (!session_init_op) {
    OpBuilder builder(&module_op.getBodyRegion());

    session_init_op = builder.create<SessionInitializerOp>(
        module_op.getLoc(), /*initializers=*/builder.getArrayAttr({}));
  }

// `module_op` is the input graph to be quantized and it should contain
// StableHLO ops.
//
// Returns a quantized `ModuleOp` in StableHLO, potentially wrapped inside a
// XlaCallModuleOp. Returns a non-OK status if quantization fails, or any of
// `saved_model_bundle` or `quantization_py_function_lib` is a nullptr.
absl::StatusOr<mlir::ModuleOp> RunQuantization(

  mlir::OpBuilder builder(&ctx);
  auto module_op = builder.create<mlir::ModuleOp>(builder.getUnknownLoc());
  // Destroy by calling destroy() to avoid memory leak since it is allocated
  // with malloc().
  const absl::Cleanup module_op_cleanup = [module_op] { module_op->destroy(); };

  const mlir::LogicalResult result = pm.run(module_op);
  EXPECT_FALSE(failed(result));

  return PopulateResultIOInfo(module_op, arg_shapes, /*use_tuple_args=*/false,
                              /*use_resource_updates_for_aliases=*/false,
                              shape_determination_fns, compilation_result);
}

static mlir::LogicalResult MlirTfToHloTextTranslateFunctionImpl(
    mlir::ModuleOp module_op, llvm::raw_ostream& output, bool via_builder) {
  if (!module_op) return mlir::failure();

namespace stablehlo::quantization {

void DisableDebugging(mlir::ModuleOp module_op) {
  module_op.walk(
      [](mlir::TF::DumpTensorOp dump_op) { dump_op.setEnabled(false); });
}

void ChangeToQuantizedFilename(mlir::ModuleOp module_op) {
  module_op.walk([](mlir::TF::DumpTensorOp dump_op) {
    dump_op.setFileName("quantized_tensor_data.pb");
  });
}