assert(op->getName().getDialectNamespace() ==
         TensorFlowDialect::getDialectNamespace());
  return tfg::InferReturnTypeComponentsForTFOp(
      location, op, op->getOperands(), graph_version, operand_as_constant_fn,
      op_result_as_shape_fn, result_element_type_fn,
      tensorflow::GetAttrValuesFromOperation, inferred_return_shapes);
}

}  // namespace TF

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

    auto tpose_arg1 = llvm::dyn_cast_or_null<TFL::TransposeOp>(
        op->getOperand(0).getDefiningOp());
    auto tpose_arg2 = llvm::dyn_cast_or_null<TFL::TransposeOp>(
        op->getOperand(1).getDefiningOp());

  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(),

           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())) {

                         "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(),

                                       decomposed_partitioned_call_callees))) {
    return failure();
  }
  if (!signature_change) return success();
  auto new_if_operands = llvm::to_vector<8>(if_op.getOperands());
  for (auto operand : if_op.getOperands()) {
    auto it = data_var_to_size_var.find(operand);
    if (it == data_var_to_size_var.end()) continue;
    new_if_operands.push_back(it->getSecond());
  }

    };
    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 &=