return LogicalResult::success(success);
  }
};

class HandleFunc : public OpRewritePattern<func::FuncOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(func::FuncOp f,
                                PatternRewriter& rewriter) const override {
    // Adjust the function type to match the block args. So this e.g. transforms
    //

    // some special cases.
    if (auto func = llvm::dyn_cast<func::FuncOp>(op)) {
      AddRegionSideEffectsForOp(func.getBody(), op);
    } else if (auto call = llvm::dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func_op = dyn_cast<func::FuncOp>(
          call.resolveCallable(&symbol_table_collection_));
      if (func_op) {

    // will be the operations inside the function body rather than representing
    // them in the function signature.
    FunctionType func_type = rewriter.getFunctionType({}, {});
    func::FuncOp func = rewriter.create<func::FuncOp>(
        loc, kImportModelDefaultGraphFuncName, func_type);
    rewriter.setInsertionPointToStart(func.addEntryBlock());
    auto executor_graph =

  RewritePatternSet patterns(ctx);
  func::FuncOp func = getOperation();

  patterns.add<AddCustomAggregationOp>(ctx, calib_opts_);
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError() << "quant-insert-custom-aggregation-ops failed.";
    signalPassFailure();
  }
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>

        .count();
  }
};

// Populates input_output_alias field in the HLO Module proto.
Status PopulateInputOutputAliasing(
    mlir::func::FuncOp main_fn,
    XlaCompiler::CompilationResult* compilation_result, bool use_tuple_args) {
  constexpr char kAliasingAttr[] = "tf.aliasing_output";
  llvm::SmallDenseMap<unsigned, unsigned> output_to_input_alias;

using mlir::OpOperand;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::success;
using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ConstOp;
using mlir::TF::FillOp;
using mlir::TF::IdentityOp;
using mlir::TF::XlaAllReduceOp;
using mlir::tf_device::ClusterOp;
using mlir::tf_device::LaunchOp;

Status HandleInputOutputArraysWithModule(
    const toco::ModelFlags& model_flags,
    mlir::OwningOpRef<mlir::ModuleOp>* module) {
  mlir::func::FuncOp entry_function = nullptr;
  for (auto func : module->get().getOps<mlir::func::FuncOp>()) {
    if (auto tf_attrs =
            func->getAttrOfType<mlir::DictionaryAttr>("tf.entry_function")) {
      // TODO(b/184697652): There could be multiple entry functions. Let's

  SymbolTable stablehlo_symbol_table(stablehlo_module);

  Builder builder(context);
  StringAttr main_func_name;
  for (auto func : stablehlo_module.getOps<func::FuncOp>()) {
    const bool is_main_func = func.getSymName() == kStablehloMainFunctionName;
    if (tf_symbol_table.lookup(func.getSymName())) {
      if (failed(stablehlo_symbol_table.renameToUnique(

namespace mlir {
namespace quant {
namespace {

using ::tensorflow::quantization::OpSet;

class PrepareLiftingPass
    : public PassWrapper<PrepareLiftingPass, OperationPass<func::FuncOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PrepareLiftingPass)

  PrepareLiftingPass() = default;

  explicit PrepareLiftingPass(OpSet op_set) { op_set_ = op_set; }

//   correct execution).
class SideEffectAnalysisInfo {
 public:
  SideEffectAnalysisInfo() = default;

  // Constructs analysis info by analyzing the given function.
  SideEffectAnalysisInfo(func::FuncOp func_op,
                         const OpSideEffectCollector& op_side_effect_collector,
                         const TF::ResourceAliasAnalysis::Info& alias_analysis,