<< "': " << error.message() << "\n";
    return 1;
  }

  // Load the MLIR module.
  mlir::DialectRegistry registry;
  registry.insert<mlir::TF::TensorFlowDialect, mlir::TFL::TensorFlowLiteDialect,
                  mlir::arith::ArithDialect, mlir::func::FuncDialect>();
  mlir::MLIRContext context(registry);

  llvm::SourceMgr source_mgr;

  void SetUp() override {
    context_ = std::make_unique<mlir::MLIRContext>();
    context_->loadDialect<mlir::arith::ArithDialect, mlir::func::FuncDialect,
                          mlir::TF::TensorFlowDialect, TensorFlowLiteDialect>();
    builder_ = std::make_unique<mlir::Builder>(context_.get());

    fused_max_unpooling_func_ =
        createMaxUnpoolingFunc(builder_.get(), {2, 4, 4, 2}, {2, 2, 2, 2});

void populateCanonicalizationPatterns(func::FuncOp func,
                                      RewritePatternSet &patterns) {
  MLIRContext *context = func.getContext();
  mlir::Dialect *tf = context->getLoadedDialect<mlir::TF::TensorFlowDialect>();
  // Load all official canonicalization patterns. Here we skip the
  // canonicalization of the ops in the tf dialect, because they couldn't
  // propagate the attributes correctly. These optimization will be played by

           "module for post-training dynamic range case";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

  void runOnOperation() override;

 private:
  Option<OpSet> target_opset_{
      *this, "target-opset", llvm::cl::init(OpSet::TF),
      llvm::cl::desc("Choose target opset."),

  if (auto cost_function =
          mlir::dyn_cast<tfrt::compiler::CostFunctionInterface>(op)) {
    cost_map_[op] = cost_function.cost();
    return;
  }

  if (!llvm::isa<mlir::TF::TensorFlowDialect>(op->getDialect())) {
    cost_map_[op] = max_arg_size_;
    return;
  }

  // Try to use its cost function if it is registered.
  const auto& registry = GetCostFunctionRegistry();

  // Constant folding is essential for the bridge. Without this check, the
  // bridge may fail with an error that is difficult to understand and not
  // actionable.
  if (!mlir::TF::TensorFlowDialect::HasConstantFoldHook()) {
    return tensorflow::errors::Internal(
        "TensorFlow dialect missing constant fold hook in TFXLA bridge phase "
        "1; this could happen if the binary doesn't link the constant fold "

    // This is a brief description of the pass.
    return "Insert quantized functions into the module";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, func::FuncDialect>();
  }

 private:
  void runOnOperation() override;

  // Returns the function library for the given quantization method and opset
  // pair.

  registry.insert<mlir::arith::ArithDialect,
                  mlir::func::FuncDialect,
                  mlir::scf::SCFDialect,
                  mlir::shape::ShapeDialect,
                  mlir::TF::TensorFlowDialect,
                  mlir::tf_device::TensorFlowDeviceDialect,
                  mlir::tf_executor::TensorFlowExecutorDialect,
                  mlir::TFR::TFRDialect>();
  // clang-format on

    // remove this restriction altogether except for the terminators that
    // require function signature change and shouldn't be
    if (op->getName().getDialectNamespace() !=
        TF::TensorFlowDialect::getDialectNamespace())
      return failure();

    // Refining types may have implications to the attached regions or symbol
    // references so do not update such ops.