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

    };
    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();

// Get total bytes transferred.
int64_t GetTransferredTensorBytes(func::CallOp from_graph,
                                  func::CallOp to_graph) {
  int64_t total_size_transferred = 0;
  for (auto input : to_graph.getOperands()) {
    Operation* input_op = input.getDefiningOp();
    if (input_op && input_op == from_graph.getOperation()) {
      auto input_type =
          mlir::dyn_cast_or_null<RankedTensorType>(input.getType());

      liveouts.snapshot_previous_state();
      return;
    }

    reverse_subgraph.insert(op);
    defined_values.insert(op->getResults().begin(), op->getResults().end());
    operands.insert(op->getOperands().begin(), op->getOperands().end());
  };

  for (Operation* op : reverse_main_func_block_ops) {
    if (!ops_to_add.contains(op)) continue;
    // When hitting a non-StableHLO op, i.e. tf.CustomAggregatorOp, start

                                       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());
  }