Operation *new_op = rewriter.create(state);
    llvm::SmallVector<Value, 4> new_results = new_op->getResults();

    // Add qint->int CastOp after output result if its original type is qint and
    // its users are not all qint->int CastOps.
    for (int i = 0; i < new_results.size(); ++i) {
      Value &result = new_results[i];
      if (IsIllegalType(result.getType()) &&

      return failure();
    }

    // Update the results.
    SmallVector<Type, 4> new_results;
    if (failed(getTypeConverter()->convertTypes(op->getResultTypes(),
                                                new_results)))
      return failure();

    // Update the regions. The dialect conversion framework wants new regions to

  Operation* return_op = body.getTerminator();
  int num_results = return_op->getNumOperands();

  // Values defined within the cluster.
  llvm::SmallVector<Value, 4> cluster_vals;
  cluster_vals.reserve(num_results);

  // New results stores values to use while replacing the old cluster op.
  llvm::SmallVector<Value, 4> new_results;
  new_results.reserve(num_results);
  for (OpOperand& operand : return_op->getOpOperands()) {

    // a tensor_list, "tfr.get_element" op is used to extract the required TF
    // op result.
    llvm::SmallVector<Value, 4> new_results;
    for (auto res : llvm::enumerate(compose_func_type.getResults())) {
      if (mlir::dyn_cast<TFRTensorType>(res.value())) {
        new_results.push_back(new_op.getResult(res.index()));
      } else if (auto list_type =
                     mlir::dyn_cast<TFRTensorListType>(res.value())) {

  SmallVector<Value, 4> new_results;
  for (auto res : llvm::enumerate(call_op.getResultTypes())) {
    Type res_type = res.value();
    if (mlir::dyn_cast<TFRTensorType>(res_type)) {
      Value new_res = new_op->getResult(res.index());
      auto casted = rewriter.create<CastOp>(loc, res_type, new_res);
      new_results.push_back(casted.getOut());
    } else if (auto list_type =

                               mlir::func::ReturnOp terminator, Value token) {
  auto new_results = llvm::to_vector(terminator.getOperands());
  new_results.push_back(token);
  builder.setInsertionPoint(terminator);
  builder.create<mlir::func::ReturnOp>(terminator.getLoc(), new_results);
  terminator.erase();
}

// Rewrites a function body and communication ops inside. Region control flow

      InsertToSymbolTable(*module, *wrap_func, func_name);
  builder.setInsertionPointAfter(result_op);
  ValueRange new_results =
      CreateFunctionCallOp(builder, call_op_loc, call_op_type,
                           new_func_name.getValue(), result_types, arguments);
  return SmallVector<Value, 4>(new_results.begin(), new_results.end());
}

SmallVector<Value, 4> LiftAsFunctionCall(OpBuilder& builder,

    // yield.
    SmallVector<Value, 4> new_results(old_num_operands);
    int next_idx = 0;
    for (int op_idx : llvm::seq<int>(0, old_num_operands))
      if (!removed_operand.test(op_idx))
        new_results[op_idx] = new_while_op.getResult(next_idx++);

    rewriter.replaceOp(while_op, new_results);
    return success();
  }
};

}  // anonymous namespace

        Value new_result = new_op_mapper.lookup(return_value);

        call_result.replaceAllUsesWith(new_result);
      }
      continue;
    }

    Operation& new_op = *builder.clone(op, arg_mapper);
    for (auto [result, new_result] :
         llvm::zip_equal(op.getResults(), new_op.getResults())) {
      new_op_mapper.map(result, new_result);
    }
  }

  call_op.erase();

          req_op.getLoc(), rewriter.getF32FloatAttr(6));
      new_result = rewriter.create<mlir::stablehlo::ClampOp>(
          req_op.getLoc(), min, new_result, max);
    } else if (func_name.contains("_relu_")) {
      auto min = rewriter.create<mlir::stablehlo::ConstantOp>(
          req_op.getLoc(), rewriter.getF32FloatAttr(0));
      new_result = rewriter.create<mlir::chlo::BroadcastMaxOp>(