func::FuncOp callee_func_op,
                                         PatternRewriter& rewriter) {
  OpBuilder::InsertionGuard insertion_guard(rewriter);

  // Create a new `CallOp` that calls `callee_func_op`.
  rewriter.setInsertionPoint(xla_call_module_op);
  auto call_op =

namespace mhlo {

// Builds body for reduce op by using the template binary op as the
// reducer op.
template <typename Op>
void BuildReduceBody(Type element_type, Region* body, OpBuilder* builder) {
  OpBuilder::InsertionGuard guard(*builder);
  Block* block = builder->createBlock(body);

  // Block arguments are scalars of the given element type.
  Type type = RankedTensorType::get(/*shape=*/{}, element_type);
  Location loc = body->getLoc();

                        {create_unknown_output_shape});

  func::FuncOp quantization_func =
      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,

                                        block.getArgumentTypes(),
                                        func_type.getResults()));

  IRMapping mapping;
  OpBuilder builder(module_op);
  OpBuilder::InsertionGuard g(builder);
  // The function has been checked to have at least one use.
  auto function_uses =
      SymbolTable::getSymbolUses(target_func, &module_op.getBodyRegion());

    if (!host_module.empty()) {
      mlir::OwningOpRef<mlir::ModuleOp> module_for_func;

      func::FuncOp func = op.GetHostFunc(&module_for_func);

      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointAfter(op->getParentOfType<func::FuncOp>());
      cloned_func = llvm::dyn_cast_or_null<func::FuncOp>(
          rewriter.clone(*func.getOperation()));

      /*controlInputs=*/control_dependencies);
  wrapper_island_op.getBody().emplaceBlock();

  // Create a NoOp inside the IslandOp.
  auto guard = OpBuilder::InsertionGuard(builder);
  builder.setInsertionPointToStart(&wrapper_island_op.GetBody());

  builder.create<TF::NoOp>(loc);
  builder.create<tf_executor::YieldOp>(loc);

  return wrapper_island_op;
}

        i, kTfSavedModelIndexPathAttr,
        ArrayAttr::get(context, {StringAttr::get(context, output_names[i])}));
  }

  // Creates PartitionedCall ops to call exported functions.
  auto guard = OpBuilder::InsertionGuard(builder);
  int arg_idx = 0;
  int result_idx = 0;
  llvm::SmallVector<Value> call_op_returns;
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (!ShouldIncludeInMainFunction(func_op)) continue;

  // Create a private function and copy all ops between arguments and results.
  auto current_func = result_op->getParentOfType<func::FuncOp>();
  auto guard = OpBuilder::InsertionGuard(builder);
  builder.setInsertionPointAfter(current_func);
  TypeRange arg_types{ValueRange{arguments}};
  TypeRange result_types{ValueRange{results}};

          "StorageCastOp must have the same quantized type");
      return failure();
    }

    // Cast to float type before the AvgPool op.
    OpBuilder::InsertionGuard g(rewriter);
    rewriter.setInsertionPointAfter(preceding_sc_op);
    auto fcast_op = rewriter.create<TF::CastOp>(
        preceding_sc_op->getLoc(), dq_arg_type.clone(rewriter.getF32Type()),

                                                 init_values_with_loop_iv);
  auto ivs_count = init_types_with_loop_iv.size();

  {
    OpBuilder::InsertionGuard guard(*builder);

    // Build up the only block in the condition region.
    Region &condition = while_op.getCond();
    Block *block = builder->createBlock(&condition);