llvm::ArrayRef<mlir::Value>{var_handle_op_in_init->getResult(0),
                                  const_op.getResult()});
}

func::FuncOp CreateSessionInitFunc(ModuleOp module) {
  constexpr char kSessionInitFuncName[] = "SessionInitializerFunction";

  mlir::OpBuilder builder(module.getBodyRegion());
  auto func_type =

    spec->quantizable_operands.insert(key);
  }
  return spec;
}

TEST_F(ApplyQuantizationParamsPropagationTest,
       ConstsUsedMultipleTimesAreDuplicated) {
  const OwningOpRef<ModuleOp> module_op_ref =
      ParseModuleOpString(kModuleTFLite);
  func::FuncOp main_fn = FindMainFuncOp(*module_op_ref);

  auto op_quant_spec_getter = [&](Operation* op) {

 public:
  AssetSinkingPass() = default;

  explicit AssetSinkingPass(llvm::StringRef saved_model_dir) {
    saved_model_dir_ = saved_model_dir.str();
  }

  void runOnOperation() override {
    ModuleOp module = getOperation();
    if (!HasTfSavedModelSemantics(module)) {
      return;
    }

    auto init_op = GetSessionInitializerOp(module);
    if (init_op == nullptr || init_op.getInitializers().empty()) {

      %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kFullyQuantizedAdd);
  ASSERT_TRUE(module_op);

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

    } else if (hlo_import_type == HloImportType::proto) {
      module = xla::HloToMlirHloTranslateFunction(content, &context, false);
    } else {
      module = mlir::OwningOpRef<mlir::ModuleOp>(
          mlir::parseSourceString<mlir::ModuleOp>(content, &context));
    }
  } else {
    // Graphdef import path.
    module = tensorflow::LoadFromGraphdefOrMlirSource(

  LogicalResult matchAndRewrite(TF::XlaCallModuleOp op,
                                PatternRewriter& rewriter) const override {
    auto module_op = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module_op);

    auto entry_func_op = dyn_cast_or_null<func::FuncOp>(
        symbol_table.lookup(GetEntryFunctionName(op)));
    if (!entry_func_op) return failure();

namespace tf2xla {
namespace v2 {

using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;

// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
tensorflow::Status RunTFXLABridge(
    ModuleOp module,
    llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,

}

void XlaRewritePass::runOnOperation() {
  ModuleOp module = getOperation();
  SymbolTable symtab(module);
  OpBuilder builder(&getContext());
  module.walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    RewriteCall(cluster_func_op, symtab, builder);
  });
}

}  // namespace

namespace TFDevice {
std::unique_ptr<OperationPass<ModuleOp>> CreateXlaRewritePass() {

 private:
  void runOnOperation() override;
};

void TestPostCalibrationComponentPass::runOnOperation() {
  ModuleOp module_op = getOperation();
  MLIRContext& ctx = getContext();

  OpPassManager pm(ModuleOp::getOperationName());

  QuantizationConfig config = QuantizationConfig::default_instance();
  config.mutable_static_range_ptq_preset();

  for (auto& op : vars) RecursiveRemove(op, erase_list, dead_ops);

  // Erases the ops.
  for (auto op : erase_list) op->erase();
}

void RemoveVariablesInSessionInitializerPass::runOnOperation() {
  ModuleOp module_op = getOperation();

  for (auto init_func_op : GetInitializerFunctions(module_op)) {
    if (!init_func_op) return;

    if (init_func_op.getBlocks().size() != 1) {