// Recreate the while op.
  OpBuilder builder(while_op);
  // Now use the filtered original operands, which will be replaced by
  // AddLoadsStoresOutsideControlFlowOp().
  auto new_while = builder.create<TF::WhileOp>(
      while_op.getLoc(), body.getFunctionType().getResults(),
      FilterRange<Value, OperandRange>(while_op.getOperands(),
                                       resource_arg_uses),

      MarkOutsideCompiled(host_if.getOperation());
    }
    if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(op)) {
      if (!HasOutsideCompilationNested(op)) return WalkResult::advance();
      OpBuilder builder(while_op);
      auto host_while = CloneEmptyWhile(while_op.getParallelIterations(),
                                        while_op.getLoc(), builder);
      const auto condition_send_recv_key =

                                TF::WhileOp* while_op) {
  // Given a while loop body function 'body_func', find the tf.While Op that
  // uses it.
  auto uses_optional = body_func.getSymbolUses(module);
  if (!uses_optional.has_value()) {
    body_func.emitOpError() << "no use of while loop body";
    return LogicalResult::failure();
  }
  *while_op = nullptr;
  for (auto& use : uses_optional.value()) {

  LogicalResult matchAndRewrite(WhileRegionOp while_op,
                                PatternRewriter &rewriter) const override {
    auto &body_block = while_op.getBody().front();
    auto &cond_block = while_op.getCond().front();
    bool changed = false;
    for (int op_idx : llvm::seq<int>(0, while_op.getNumOperands())) {
      auto body_arg = body_block.getArgument(op_idx);

            if (!func) continue;
            next_values_to_visit.push_back(
                func.getArgument(next_op_operand_num));
          } else if (auto while_op =
                         llvm::dyn_cast_or_null<TF::WhileOp>(next_op)) {
            func::FuncOp func = while_op.body_function();
            auto func_argument = func.getArgument(next_op_operand_num);
            // Check if the op is returned without mutation. Returning values

          return logical_device;
        }

        if (auto while_op = llvm::dyn_cast<mlir::TF::WhileRegionOp>(owner)) {
          const int operand_number = use.getOperandNumber();
          next_values_to_visit.push_back(
              while_op.getCond().front().getArgument(operand_number));
          next_values_to_visit.push_back(
              while_op.getBody().front().getArgument(operand_number));
          continue;

          while_op.getInput().getTypes(), while_op.getOutput().getTypes(),
          while_op.getBody().getArgumentTypes());
      return PropagateShapeToRegions(compatible_types,
                                     {&while_op.getCond(), &while_op.getBody()},
                                     max_iterations);
    }
    return PropagateShapeToRegions(while_op.getInput().getTypes(),

      result =
          CanonicalizeFunctionalIfCase(case_op, branches, case_op.getInput());
    } else if (auto while_op = dyn_cast<TF::WhileOp>(op)) {
      if (while_op.cond_function().walk(check_while_cond).wasInterrupted())
        return WalkResult::interrupt();
      result = CanonicalizeFunctionalWhile(while_op);
    } else if (isa<TF::IfRegionOp, TF::CaseRegionOp, tf_device::ClusterOp>(
                   op)) {

        if (failed(LowerIfOp(if_op))) {
          return signalPassFailure();
        }
        break;
      }
      if (WhileOp while_op = llvm::dyn_cast<WhileOp>(op)) {
        if (failed(LowerWhileOp(while_op))) {
          return signalPassFailure();
        }
        break;
      }
    }
  }
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>>

    Value collection, ModuleOp module,
    llvm::function_ref<std::optional<Type>(Operation*)> infer_from_op) {
  for (auto& use : collection.getUses()) {
    if (auto while_op = llvm::dyn_cast<TF::WhileOp>(use.getOwner())) {
      auto body = while_op.body_function();
      assert(body);
      auto type_from_body = GetElementTypeFromAccess(
          body.getArgument(use.getOperandNumber()), module, infer_from_op);