// Set up the 'then' block.
  Block* then_block = builder.createBlock(merge_block);
  Operation* call_op = CallFn(loc, get_operand, op.then_function(), &builder);

  auto get_then_result = [&](int i) { return call_op->getResult(i); };
  JumpToBlock(loc, get_then_result, merge_block, &builder);

  // Set up the 'else' block.
  Block* else_block = builder.createBlock(merge_block);

void PopulateDatasetWhileCond(OpBuilder builder, WhileRegionOp dataset_while,
                              Location loc) {
  auto& cond_region = dataset_while.getCond();
  Block* cond_block = builder.createBlock(&cond_region);
  auto while_input_types = dataset_while.getOperandTypes();
  cond_block->addArguments(
      while_input_types, SmallVector<Location>(while_input_types.size(), loc));

  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()));
  func_input_arg = block->getArgument(0);
  return quantization_func;
}

                   bool forward_block_args) {
  assert(caller_region.empty() &&
         "Expected empty region for newly created ops");
  OpBuilder builder(caller_region);
  Block* entry = builder.createBlock(&caller_region);

  auto loc = op->getLoc();
  if (use_region_args) {
    auto inputs = func.getFunctionType().getInputs();
    entry->addArguments(inputs, SmallVector<Location>(inputs.size(), loc));

                                num_loop_carried, extern_values, types, loc);
  OpBuilder b(region);
  // The body of the region is empty/has been outlined into the function.
  auto block = b.createBlock(&region);
  SmallVector<Value, 4> new_operands;
  new_operands.reserve(types.size());
  for (Type t : llvm::ArrayRef(types).drop_back(extern_values.size()))
    new_operands.push_back(block->addArgument(t, loc));

  const auto func_type = builder.getFunctionType(/*inputs=*/{}, /*results=*/{});

  auto init_func = builder.create<func::FuncOp>(
      loc, /*sym_name=*/"init_func_restore_op", func_type);
  builder.createBlock(&init_func.getBody(), /*insertPt=*/init_func.begin(),
                      /*arg_types=*/{}, /*arg_locs=*/{});

  init_func->setAttr(kTfSavedModelExportedNamesAttr,
                     builder.getStrArrayAttr(

  auto func_type = FunctionType::get(context, arg_types, result_types);
  auto main_func = builder.create<func::FuncOp>(
      module_op.getLoc(), kImportModelDefaultGraphFuncName, func_type);
  builder.createBlock(&main_func.getBody(), main_func.begin(), arg_types,
                      arg_locs);
  SmallVector<NamedAttribute> func_attrs;
  func_attrs.push_back(
      {StringAttr::get(context, "inputs"),

  if (call_op_type == FunctionCallOpType::TFXlaCallModuleOp) {
    wrap_func->setAttr(TF::kFromXlaCallModuleAttrName, builder.getUnitAttr());
  }
  wrap_func->setAttr(kFusedFunctionAttr, builder.getUnitAttr());
  builder.createBlock(&wrap_func.getBody(), wrap_func.begin(), arg_types,
                      arg_locs);

  IRMapping mapping;
  for (int32_t i : llvm::seq<int32_t>(0, arguments.size())) {

  stablehlo_func_op.setVisibility(SymbolTable::Visibility::Private);
  stablehlo_func_op->setAttr(TF::kFromXlaCallModuleAttrName,
                             builder.getUnitAttr());

  builder.createBlock(&stablehlo_func_op.getBody(), stablehlo_func_op.begin(),
                      arg_types, arg_locs);

  IRMapping mapper;
  // stablehlo_func_op has 1 extra arg for platform index.

  auto cluster_op =
      builder.create<tf_device::ClusterOp>(loc, return_types, policy);

  // Create block in cluster_op's region and move 'cluster.operations' into
  // it.
  auto block = builder.createBlock(&cluster_op.getBody());
  auto block_end = block->end();
  for (auto op : cluster.operations) op->moveBefore(block, block_end);

  // Add 'tf_device::ReturnOp' at the end of the block.