def __init__(self, *args, **kwargs) -> None: ...

class Block_Iterator:
    def __init__(self, *args, **kwargs) -> None: ...

class Builder:
    def __init__(self, arg0: MLIRContext) -> None: ...
    def getFunctionType(self, arg0: list[Type], arg1: list[Type]) -> FunctionType: ...

class FloatType(Type):
    def __init__(self, *args, **kwargs) -> None: ...
    def getBF16(self) -> FloatType: ...
    def getF16(self) -> FloatType: ...

      rewriter.create<func::FuncOp>(input_op->getLoc(), func_name, func_type);

  OpBuilder::InsertionGuard guard = OpBuilder::InsertionGuard(rewriter);
  ArrayRef<Type> inputs = quantization_func.getFunctionType().getInputs();
  Block* block = rewriter.createBlock(
      &quantization_func.getBody(), quantization_func.begin(), inputs,
      SmallVector<Location>(inputs.size(), quantization_func.getLoc()));

  main_func.getArguments().front().setType(tensor_string_type);
  main_func.setType(
      FunctionType::get(builder.getContext(), {tensor_string_type},
                        main_func.getFunctionType().getResults()));
  // Name of the main function for the eventual executable needs to be set.
  main_func->setAttr(kTfSavedModelExportedNamesAttr,
                     builder.getStrArrayAttr({func_name}));

      .def("getArguments",
           [](mlir::func::FuncOp& f) { return f.getArguments().vec(); })
      .def("getName", [](mlir::func::FuncOp& f) { return f.getName().str(); })
      .def("getType", &mlir::func::FuncOp::getFunctionType);

  py::class_<mlir::func::ReturnOp>(m, "ReturnOp")
      .def("create",
           [](mlir::OpBuilder& opb, mlir::Location loc,
              std::vector<mlir::Value> values) -> mlir::Operation* {

    // Update the function type.
    func.setType(mlir::FunctionType::get(module.getContext(),
                                         func.getBody().getArgumentTypes(),
                                         func.getFunctionType().getResults()));
  }
  return success();
}

}  // namespace tf_saved_model

  DCOMMENT_OP(op.getOperation(), "Infer shape for if ");
  bool changed = false;
  auto then_results =
      op.ResolveThenFunction(&symbol_table_).getFunctionType().getResults();
  auto else_results =
      op.ResolveElseFunction(&symbol_table_).getFunctionType().getResults();
  for (auto it : llvm::zip(op.getResults(), then_results, else_results)) {
    // If then and else types do not match, skip refinement for that result.

    mlir::Region *getCallableRegion();

    /// Returns the argument types of this function.
    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }

    /// Returns the result types of this function.
    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }

    // Get the names of all defined attributes, including both derived and
    // non-derived ones.

    builder.setInsertionPointToStart(module_->getBody());
    auto func = builder.create<mlir::func::FuncOp>(  //
        builder.getUnknownLoc(), "main", builder.getFunctionType({}, {}));
    func->setAttr("tfl.func", builder.getUnitAttr());
    builder.setInsertionPointToStart(func.addEntryBlock());
    llvm::SmallVector<int> shape{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

    for (auto ret_op : ret->getOperands()) {
      return_types.push_back(ret_op.getType());
    }
    auto newType =
        FunctionType::get(rewriter.getContext(), argument_types, return_types);
    if (f.getFunctionType() == newType) {
      return failure();
    }
    rewriter.modifyOpInPlace(f, [&] { f.setType(newType); });

    // Adjust the type of the return values callers of the function to

  llvm::StringMap<int> shared_name_to_arg_idx;
  llvm::SmallVector<std::pair<Operation*, int>> lifted_op_and_arg_idx;
  Block& block = target_func.front();
  auto func_type = target_func.getFunctionType();

  for (Operation& op : block.without_terminator()) {
    StringRef shared_name = GetSharedName(&op);
    if (shared_name.empty() || !IsHashTableOp(&op)) continue;