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] :
           llvm::zip_equal(call_op.getResults(), op.getOperands())) {

  const int operands_size = while_op.getNumOperands() + extra_operands.size();
  SmallVector<Value, 4> operands;
  operands.reserve(operands_size);
  operands.append(while_op.getOperands().begin(), while_op.getOperands().end());
  operands.append(extra_operands.begin(), extra_operands.end());
  SmallVector<Type, 4> new_types;
  new_types.reserve(operands_size);
  new_types.append(while_op.getResultTypes().begin(),

LogicalResult FoldOperation(OpBuilder& builder, Operation* op,
                            SmallVector<Value>& results) {
  SmallVector<ElementsAttr> inputs;
  for (auto operand : op->getOperands()) {
    auto preceding_const_op = operand.getDefiningOp<TF::ConstOp>();
    if (preceding_const_op) {
      inputs.push_back(preceding_const_op.getValue());
      continue;
    }

                         "single block\n";
      return;
    }
    Block &block = body.front();
    for (const auto &arg_mapping :
         llvm::zip(block.getArguments(), op->getOperands())) {
      std::get<0>(arg_mapping).replaceAllUsesWith(std::get<1>(arg_mapping));
    }
    op->getBlock()->getOperations().splice(op->getIterator(),

    };
    if (!llvm::all_of(operation->getOperands(), foldable_opr)) {
      return std::nullopt;
    }
    return FoldAdaptor(operation);
  }

  // Gets a list of ElementsAttr behind each constant operand.
  llvm::SmallVector<ElementsAttr> OperandData() {
    llvm::SmallVector<ElementsAttr> res;
    res.reserve(operation_->getNumOperands());
    for (auto opr : operation_->getOperands()) {

    if (op->hasAttr(kShardingAttr) && op.getCallTargetName() == kShardingName &&
        op.getNumOperands() == 1 && op.getNumResults() == 1 &&
        op.getOperands().front().getType() ==
            op.getResults().front().getType()) {
      rewriter.replaceOp(op, op.getOperands());
      return success();
    }
    return failure();
  }
};

namespace {

bool IsShloMainFuncOp(func::FuncOp func_op) {

    if (common_inner_operands.empty()) {
      common_inner_operands.insert(common_inner_operands.end(),
                                   op->GetOperands().begin(),
                                   op->GetOperands().end());
    } else {
      common_inner_operands.clear();
      absl::c_copy_if(op->GetOperands(),
                      std::back_inserter(common_inner_operands),
                      [&](Predicate* sub_op) {

                                PatternRewriter &rewriter) const final {
    Operation *op = srcop.getOperation();
    bool allTypesFp = true;
    bool allTypesQuantizedOrInt = true;
    for (auto operand : op->getOperands()) {
      ShapedType type = mlir::dyn_cast<ShapedType>(operand.getType());
      if (!type) continue;
      allTypesFp &= !mlir::isa<quant::QuantizedType>(type.getElementType());
      allTypesQuantizedOrInt &=

        composite->getLoc(), composite->getResultTypes(),
        composite->getOperands().slice(2, 3), composite.getName(),
        CustomOption(&builder, custom_option_buffer));
  }
  return builder.create<TFL::CustomOp>(
      composite->getLoc(), composite->getResultTypes(),
      composite->getOperands(), composite.getName(),
      CustomOption(&builder, custom_option_buffer));
}

}  // namespace

    }

    const bool any_blargs = llvm::any_of(
        op->getOperands(),
        [](Value opr) -> bool { return llvm::isa<BlockArgument>(opr); });
    if (any_blargs) {
      return failure();
    }

    TFL::TransposeOp tpose_arg = nullptr;
    Operation *cst_arg = nullptr;

    for (Value opr : op->getOperands()) {
      auto *defining_op = opr.getDefiningOp();