SmallVector<Location> arg_locs;
  for (Value arg : arguments) {
    arg_locs.push_back(arg.getLoc());
  }

  auto wrap_func = builder.create<func::FuncOp>(location, func_name, func_type);
  wrap_func.setVisibility(SymbolTable::Visibility::Private);
  // The callee function for TF::XlaCallModuleOp must have this attribute.
  if (call_op_type == FunctionCallOpType::TFXlaCallModuleOp) {

  auto stablehlo_func_op = builder.create<func::FuncOp>(
      module_op.getLoc(), CreateStablehloFunctionName(stablehlo_func_id),
      FunctionType::get(ctx, arg_types, result_types));
  stablehlo_func_op.setVisibility(SymbolTable::Visibility::Private);
  stablehlo_func_op->setAttr(TF::kFromXlaCallModuleAttrName,
                             builder.getUnitAttr());

    auto then_branch_op =
        rewriter.create<func::FuncOp>(loc, "cond_true", func_type);
    CreateCondTrueBranch(op, shape_dtype, result_type, then_branch_op,
                         &rewriter);
    then_branch_op.setVisibility(func::FuncOp::Visibility::Private);

    // Constructs `else_branch`, which is executed when `if_cond` evaluates to
    // false.
    auto else_branch_op =

    auto visibility = function.getName() == graph_func_name
                          ? mlir::func::FuncOp::Visibility::Public
                          : mlir::func::FuncOp::Visibility::Private;
    function.setVisibility(visibility);
  }
  VLOG(2) << "Imported: "
          << tensorflow::DumpMlirOpToFile("tf_mlir_imported_base",
                                          module.get());
  return module;
}