new_while_operands.push_back(it->getSecond());
  }
  auto new_while = builder.create<TF::WhileOp>(
      while_op.getLoc(), body.getFunctionType().getInputs(), new_while_operands,
      while_op->getAttrs());
  for (int64_t i = 0; i < while_op.getNumResults(); ++i) {
    if (!mlir::isa<TF::ResourceType>(
            getElementTypeOrSelf(while_op.getOperand(i).getType()))) {
      continue;
    }

  auto new_cluster = builder->create<mlir::tf_device::ClusterOp>(
      cluster.getLoc(), new_cluster_result_types,
      /*operands=*/llvm::ArrayRef<Value>{}, cluster->getAttrs());
  new_cluster.getBody().takeBody(cluster.getBody());

  auto operand_not_in_cluster = [&](mlir::OpOperand& operand) {
    return !new_cluster.getOperation()->isProperAncestor(operand.getOwner());
  };

  }

  // Rebuild the new operation with lesser number of results.
  OpBuilder builder(op);
  Operation *new_op = Operation::create(
      op->getLoc(), op->getName(), new_result_types, op->getOperands(),
      op->getAttrs(), op->getPropertiesStorage(), op->getSuccessors(),
      op->getNumRegions());
  builder.insert(new_op);

  // Move region bodies to the new operation.

      OperationState new_state(candidate_op->getLoc(),
                               candidate_op->getName().getStringRef(), inputs,
                               output_types, candidate_op->getAttrs());
      for (int i = 0; i < candidate_op->getNumRegions(); ++i) {
        new_state.addRegion();
      }
      Operation* quantized_op = rewriter.create(new_state);
      if (candidate_op->getNumRegions() != 0) {

      new_result_types.push_back(while_op.getResultTypes()[i]);
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {
    if (!while_op.getResult(i).use_empty()) {

  for (auto extra_operand : extra_operands)
    new_types.push_back(extra_operand.getType());

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_types, operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());
  while_op.replaceAllUsesWith(
      new_while_op.getResults().take_front(while_op.getNumResults()));

  // Skip derived `operandSegmentSizes` attribute as custom print format of
  // operands holds enough information to calculate these variadic operand list
  // lengths.
  p.printOptionalAttrDict(
      getOperation()->getAttrs(),
      /*elidedAttrs=*/ArrayRef<StringRef>{kOperandSegmentSizesAttr});
  p << ' ';
  p.printRegion(getBody(), /*printEntryBlockArgs=*/false);
}

namespace {

void SetFunctionPrivate(func::FuncOp func) {
  func.setVisibility(SymbolTable::Visibility::Private);

  // The `tf_saved_model` attributes can only be applied to public functions.
  for (auto& attr : func->getAttrs()) {
    StringRef attr_name = attr.getName().getValue();
    if (attr_name.starts_with("tf_saved_model.")) {
      func->removeAttr(attr_name);
    }
  }

    // contraction, with two additions: the list of ops which have been fused
    // together; epsilon (only with FusedBatchNorm).
    std::vector<NamedAttribute> attrs = contraction->getAttrs();
    ArrayAttr fused_ops_attr = ArrayAttr::get(context, fused_ops);
    attrs.push_back(
        NamedAttribute(StringAttr::get(context, "fused_ops"), fused_ops_attr));

    if (failed(max_unpooling.RewriteFunc())) {
      return signalPassFailure();
    }
  } else {
    // We will look for the `tfl_fusable_op` attribute and fuse as a custom op.
    DictionaryAttr dict_attr = attr.getAttrs();

    SmallVector<std::pair<StringRef, Attribute>, 4> attributes;
    bool tfl_fusable_op = false;
    for (auto attr_item : dict_attr) {
      // Push other attributes except the TFLFusableOp.