func_op.getArguments().take_back(call_op.getNumOperands()),
           call_op.getOperands())) {
    arg_mapper.map(func_arg, operand);
  }

  Region& function_body = func_op.getBody();
  IRMapping new_op_mapper;
  for (Operation& op : function_body.getOps()) {
    if (llvm::isa<func::ReturnOp>(op)) {
      for (auto [call_result, return_value] :

  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

  op->replaceAllUsesWith(new_op.getResults());
  op->erase();
}

void LegalizeWhilePass::RunOnFunction(func::FuncOp func) {
  // Convert all TF WhileOps inside the function body to TFL While ops.
  func.getBody().walk([](TF::WhileOp while_op) { RunOnWhile(while_op); });
}

  // Unwrap from any ControlNodeOps or CustomTfOp.
  if (auto control_node = dyn_cast<mlir::TFL::ControlNodeOp>(*op)) {
    return IsFlexDelegate(&control_node.getBody().front().front());
  }
  if (auto custom_tf_op = dyn_cast<mlir::TFL::CustomTfOp>(*op)) {
    return IsFlexDelegate(&custom_tf_op.getBody().front().front());
  }

  // Our MLIR might be the result of a conversion from a previously generated
  // flatbuffer file.

  getOperation().walk(
      [&](tf_device::ClusterOp cluster) { clusters.push_back(cluster); });

  for (tf_device::ClusterOp cluster : clusters) {
    RemoveIdentityFromRegion(cluster.getBody());
    auto reachable_funcs = CollectReachableFunctions(cluster.getBody());
    for (func::FuncOp reachable_func : reachable_funcs)
      RemoveIdentityFromRegion(*reachable_func.getCallableRegion());
  }
}

}  // anonymous namespace

      if (block_arg && block_arg.getOwner() == &replicate.GetBody()) {
        // Only mirrored args of ReplicateOp can be annotated.
        if (mirrored_replicate_args.count(block_arg.getArgNumber()) == 0) {
          continue;
        }
      } else if (!operand.getParentRegion()->isProperAncestor(
                     &replicate.getBody())) {
        // Not a replication-invariant operand.
        continue;
      }

    return;
  }
  SymbolTable symbol_table(module);

  for (auto func : module.getOps<func::FuncOp>()) {
    llvm::BitVector args_to_erase(func.getNumArguments());
    OpBuilder builder(func.getBody());

    for (int i = 0, e = func.getNumArguments(); i < e; ++i) {
      SmallVector<TF::InitializeTableFromTextFileV2Op, 4>
          init_table_from_text_file_ops_to_erase;

    GraphOp graph, TF::OpToParallelIdsMap& op_to_parallel_ids_map) {
  for (Operation& op : graph.GetBody()) {
    auto island = dyn_cast<IslandOp>(&op);
    if (!island) continue;

    // We call `VerifyExportSuitable` in the beginning of the pass, so every
    // island wraps a single op.
    Operation& wrapped_op = island.GetBody().front();
    TF::ParallelExecutionIdPairs id_pairs;

  Region* cluster_region = &cluster.getBody();
  llvm::SmallSetVector<Operation*, 4> head_outside_compiled_ops;

  // Traverse the graph in topological order to find all outside compiled ops
  // at head of TPU computation or unary ops that are only used by other outside
  // compiled ops.
  auto cluster_ops = cluster.GetBody().without_terminator();
  for (Operation& cluster_op : cluster_ops) {

          GetOpOfValue(cluster_func_operand.value()));
      if (!device_launch_op) continue;
      for (auto result : llvm::zip(
               device_launch_op.getResults(),
               device_launch_op.GetBody().getTerminator()->getOperands())) {
        if (std::get<0>(result) == cluster_func_operand.value() &&
            llvm::isa<TF::TPUAnnotateTensorsWithDynamicShapeOp>(
                std::get<1>(result).getDefiningOp())) {

        super(name, parameters);
        this.parameterToReturn = parameterToReturn;
    }

    public int getParameterToReturn() {
        return parameterToReturn;
    }

    @Override
    protected String getBody() {
        return "=> return parameter " + parameterToReturn;
    }

    @Override
    public Expression evaluate(List<Expression> arguments) {
        return arguments.get(parameterToReturn).asMacroExpansion();