if (shared_name.empty()) return false;

  for (func::FuncOp init_func_op :
       tf_saved_model::GetInitializerFunctions(module_op)) {
    for (Operation& op : init_func_op.getBody().getOps()) {
      StringRef other_shared_name = GetSharedName(&op);
      if (IsHashTableOp(&op) && other_shared_name == shared_name) {
        return true;
      }
    }
  }
  return false;
}

  mlir::Builder builder(context);
  StringAttr resource_name_id = builder.getStringAttr(kResourceNameArgAttr);

  SmallSet<StringRef, 4> resource_names;

  for (func::FuncOp func : module.getOps<func::FuncOp>()) {
    for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
      auto resource_arg =
          func.getArgAttrOfType<StringAttr>(i, kResourceNameArgAttr);
      if (!resource_arg) continue;

    if (!main_func) continue;

    SymbolTable symbol_table(module_op);
    for (auto call_op : main_func.getOps<TF::PartitionedCallOp>()) {
      func_ops.push_back(dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(
          mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue())));
    }
    for (auto call_op : main_func.getOps<TF::StatefulPartitionedCallOp>()) {
      func_ops.push_back(

  bool need_prefix_for_output_name = false;
  std::vector<StringRef> fn_input_name_vec, fn_output_name_vec;
  StringSet<> input_name_set, output_name_set;
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (!ShouldIncludeInMainFunction(func_op)) continue;
    if (auto tf_attrs =
            func_op->getAttrOfType<DictionaryAttr>(kEntryFunctionAttr)) {

template <typename GemmStyleOp>
// Match for all gemm_style op and check for possible fusions.
LogicalResult MatchGemmStyleOp(func::FuncOp entry_func_op) {
  const auto op_iterator_range = entry_func_op.getOps<GemmStyleOp>();
  if (op_iterator_range.empty()) {
    LLVM_DEBUG(llvm::dbgs() << "Function does not have "
                            << GemmStyleOp::getOperationName() << " op.\n");
    return failure();
  }

  // to store its `SymbolTable` to `SymbolTableCollection`.
  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(

    builder.create<YieldOp>(island_op.getLoc(), yield_operands);
  }

  // Outline all the transitively called functions by moving them in the
  // outlined module.
  for (func::FuncOp func : outlined_module.getOps<func::FuncOp>()) {
    func.walk([&](Operation *op) {
      for (NamedAttribute attr : op->getAttrs()) {
        if (auto symbol_ref =
                mlir::dyn_cast<FlatSymbolRefAttr>(attr.getValue())) {

  LogicalResult UnifyReplicationInfo(func::FuncOp func) {
    auto new_repl_info =
        builder.getStringAttr(func.getSymName().str() + "_repl_info");
    for (auto& op : func.getRegion().getOps()) {
      if (auto compile_op = llvm::dyn_cast<TF::TPUCompilationResultOp>(op)) {
        UpdateReplicationAttr(compile_op, new_repl_info);
      } else {
        UpdateReplicationAttr(&op, new_repl_info);
      }

  FailureOr<OpT> FindFirstOpFromMainFunc(ModuleOp module_op) {
    func::FuncOp main_func_op = FindMainFuncOp(module_op);
    if (main_func_op == nullptr) return failure();

    auto ops = main_func_op.getOps<OpT>();
    if (ops.empty()) return failure();

    return *ops.begin();
  }

  std::unique_ptr<MLIRContext> ctx_;
  OpBuilder builder_;
};

}  // namespace mlir::quant

    if (function_name.contains("conv")) {
      return (*(func.getOps<mlir::stablehlo::ConvolutionOp>().begin()))
          .getDimensionNumbers()
          .getKernelOutputFeatureDimension();
    } else if (function_name.contains("dot_general")) {
      auto dot = *(func.getOps<mlir::stablehlo::DotGeneralOp>().begin());
      const ::mlir::stablehlo::DotDimensionNumbersAttr dimension_numbers =